ob-metaflow-extensions 1.1.162rc0__py2.py3-none-any.whl → 1.1.162rc1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/ollama/__init__.py

@@ -2,8 +2,10 @@ from metaflow.decorators import StepDecorator
 from metaflow import current
 import functools
 import os
+import threading
 
-from .ollama import OllamaManager
+from .ollama import OllamaManager, OllamaRequestInterceptor
+from .status_card import OllamaStatusCard
 from ..card_utilities.injector import CardDecoratorInjector
 
 __mf_promote_submodules__ = ["plugins.ollama"]
@@ -38,10 +40,16 @@ class OllamaDecorator(StepDecorator, CardDecoratorInjector):
         Determines where and how to run the Ollama process.
     force_pull: bool
         Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
-    skip_push_check: bool
-        Whether to skip the check that populates/overwrites remote cache on terminating an ollama model.
+    cache_update_policy: str
+        Cache update policy: "auto", "force", or "never".
+    force_cache_update: bool
+        Simple override for "force" cache update policy.
     debug: bool
         Whether to turn on verbose debugging logs.
+    circuit_breaker_config: dict
+        Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
+    timeout_config: dict
+        Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
     """
 
     name = "ollama"
@@ -49,8 +57,22 @@ class OllamaDecorator(StepDecorator, CardDecoratorInjector):
         "models": [],
         "backend": "local",
         "force_pull": False,
-        "skip_push_check": False,
+        "cache_update_policy": "auto",  # "auto", "force", "never"
+        "force_cache_update": False,  # Simple override for "force"
         "debug": False,
+        "circuit_breaker_config": {
+            "failure_threshold": 3,
+            "recovery_timeout": 60,
+            "reset_timeout": 30,
+        },
+        "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
+        },
+        "card_refresh_interval": 10,  # seconds - how often to update the status card
     }
 
     def step_init(
@@ -61,29 +83,143 @@ class OllamaDecorator(StepDecorator, CardDecoratorInjector):
         )
         self.flow_datastore_backend = flow_datastore._storage_impl
 
+        # Attach the ollama status card
+        self.attach_card_decorator(
+            flow,
+            step_name,
+            "ollama_status",
+            "blank",
+            refresh_interval=self.attributes["card_refresh_interval"],
+        )
+
     def task_decorate(
         self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
     ):
         @functools.wraps(step_func)
         def ollama_wrapper():
+            self.ollama_manager = None
+            self.request_interceptor = None
+            self.status_card = None
+            self.card_monitor_thread = None
+
             try:
+                # Initialize status card and monitoring
+                self.status_card = OllamaStatusCard(
+                    refresh_interval=self.attributes["card_refresh_interval"]
+                )
+
+                # Start card monitoring in background
+                def monitor_card():
+                    try:
+                        self.status_card.on_startup(current.card["ollama_status"])
+
+                        while not getattr(
+                            self.card_monitor_thread, "_stop_event", False
+                        ):
+                            try:
+                                # Trigger card update with current data
+                                self.status_card.on_update(
+                                    current.card["ollama_status"], None
+                                )
+                                import time
+
+                                time.sleep(self.attributes["card_refresh_interval"])
+                            except Exception as e:
+                                if self.attributes["debug"]:
+                                    print(f"[@ollama] Card monitoring error: {e}")
+                                break
+                    except Exception as e:
+                        if self.attributes["debug"]:
+                            print(f"[@ollama] Card monitor thread error: {e}")
+                        self.status_card.on_error(current.card["ollama_status"], str(e))
+
+                self.card_monitor_thread = threading.Thread(
+                    target=monitor_card, daemon=True
+                )
+                self.card_monitor_thread._stop_event = False
+                self.card_monitor_thread.start()
+
+                # Initialize OllamaManager with status card
                 self.ollama_manager = OllamaManager(
                     models=self.attributes["models"],
                     backend=self.attributes["backend"],
                     flow_datastore_backend=self.flow_datastore_backend,
                     force_pull=self.attributes["force_pull"],
-                    skip_push_check=self.attributes["skip_push_check"],
+                    cache_update_policy=self.attributes["cache_update_policy"],
+                    force_cache_update=self.attributes["force_cache_update"],
                     debug=self.attributes["debug"],
+                    circuit_breaker_config=self.attributes["circuit_breaker_config"],
+                    timeout_config=self.attributes["timeout_config"],
+                    status_card=self.status_card,
                 )
+
+                # Install request protection by monkey-patching ollama package
+                self.request_interceptor = OllamaRequestInterceptor(
+                    self.ollama_manager.circuit_breaker, self.attributes["debug"]
+                )
+                self.request_interceptor.install_protection()
+
+                if self.attributes["debug"]:
+                    print(
+                        "[@ollama] OllamaManager initialized and request protection installed"
+                    )
+
             except Exception as e:
+                if self.status_card:
+                    self.status_card.add_event(
+                        "error", f"Initialization failed: {str(e)}"
+                    )
+                    try:
+                        self.status_card.on_error(current.card["ollama_status"], str(e))
+                    except:
+                        pass
                 print(f"[@ollama] Error initializing OllamaManager: {e}")
                 raise
+
             try:
+                if self.status_card:
+                    self.status_card.add_event("info", "Starting user step function")
                 step_func()
+                if self.status_card:
+                    self.status_card.add_event(
+                        "success", "User step function completed successfully"
+                    )
             finally:
-                self.ollama_manager.terminate_models()
-            if self.attributes["debug"]:
-                print(f"[@ollama] process statuses: {self.ollama_manager.processes}")
-                print(f"[@ollama] process runtime stats: {self.ollama_manager.stats}")
+                # Remove request protection first (before terminating models)
+                if self.request_interceptor:
+                    self.request_interceptor.remove_protection()
+                    if self.attributes["debug"]:
+                        print("[@ollama] Request protection removed")
+
+                # Then cleanup ollama manager (while card monitoring is still active)
+                if self.ollama_manager:
+                    self.ollama_manager.terminate_models()
+
+                # Give the card a moment to render the final shutdown events
+                if self.card_monitor_thread and self.status_card:
+                    import time
+
+                    # Trigger one final card update to capture all shutdown events
+                    try:
+                        self.status_card.on_update(current.card["ollama_status"], None)
+                    except Exception as e:
+                        if self.attributes["debug"]:
+                            print(f"[@ollama] Final card update error: {e}")
+                    time.sleep(2)  # Allow final events to be rendered
+
+                # Now stop card monitoring
+                if self.card_monitor_thread:
+                    self.card_monitor_thread._stop_event = True
+
+                if self.ollama_manager and self.attributes["debug"]:
+                    print(
+                        f"[@ollama] process statuses: {self.ollama_manager.processes}"
+                    )
+                    print(
+                        f"[@ollama] process runtime stats: {self.ollama_manager.stats}"
+                    )
+                    print(
+                        f"[@ollama] Circuit Breaker status: {self.ollama_manager.circuit_breaker.get_status()}"
+                    )
 
         return ollama_wrapper