ob-metaflow-extensions 1.1.162rc0__tar.gz → 1.1.163__tar.gz

{ob-metaflow-extensions-1.1.162rc0/ob_metaflow_extensions.egg-info → ob_metaflow_extensions-1.1.163}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
-Name: ob-metaflow-extensions
-Version: 1.1.162rc0
+Name: ob_metaflow_extensions
+Version: 1.1.163
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob-metaflow-extensions-1.1.162rc0 → ob_metaflow_extensions-1.1.163}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -41,6 +41,9 @@ def get_boto3_session(role_arn=None, session_vars=None):
     import boto3
     import botocore
     from metaflow_extensions.outerbounds.plugins.auth_server import get_token
+    from metaflow_extensions.outerbounds.plugins.aws.assume_role import (
+        OBP_ASSUME_ROLE_ARN_ENV_VAR,
+    )
 
     from hashlib import sha256
     from metaflow.util import get_username
@@ -49,6 +52,10 @@ def get_boto3_session(role_arn=None, session_vars=None):
 
     token_info = get_token("/generate/aws")
 
+    # Check if the assume_role decorator has set a role ARN via environment variable
+    # This takes precedence over CSPR role
+    decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+
     # Write token to a file. The file name is derived from the user name
     # so it works with multiple users on the same machine.
     #
@@ -69,12 +76,13 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("cspr_role_arn"):
         cspr_role = token_info["cspr_role_arn"]
 
-    if cspr_role:
-        # If CSPR role is set, we set it as the default role to assume
-        # for the AWS SDK. We do this by writing an AWS config file
+    # If assume_role decorator is used, prioritize it over CSPR role
+    effective_role = decorator_role_arn or cspr_role
+    if effective_role:
+        # If we have either a decorator role or CSPR role, set up AWS config
         # with two profiles. One to get credentials for the task role
         # in exchange for the OIDC token, and second to assume the
-        # CSPR role using the task role credentials.
+        # effective role using the task role credentials.
         import configparser
         from io import StringIO
 
@@ -86,9 +94,9 @@ def get_boto3_session(role_arn=None, session_vars=None):
             "web_identity_token_file": token_file,
         }
 
-        # CSPR role profile (default)
-        aws_config["profile cspr"] = {
-            "role_arn": cspr_role,
+        # Effective role profile (decorator role or CSPR role)
+        aws_config["profile effective"] = {
+            "role_arn": effective_role,
             "source_profile": "task",
         }
 
@@ -104,7 +112,7 @@ def get_boto3_session(role_arn=None, session_vars=None):
             tmp_aws_config_file = f.name
         os.rename(tmp_aws_config_file, aws_config_file)
         os.environ["AWS_CONFIG_FILE"] = aws_config_file
-        os.environ["AWS_PROFILE"] = "cspr"
+        os.environ["AWS_PROFILE"] = "effective"
     else:
         os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
         os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
@@ -116,21 +124,22 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("region"):
         os.environ["AWS_DEFAULT_REGION"] = token_info["region"]
 
-    if cspr_role:
-        #  The generated AWS config will be used here since we set the
-        #  AWS_CONFIG_FILE environment variable above.
-        if role_arn == USE_CSPR_ROLE_ARN_IF_SET:
-            # Otherwise start from the default profile, assuming CSPR role
-            session = boto3.session.Session(profile_name="cspr")
+    if effective_role:
+        if role_arn == decorator_role_arn or role_arn == USE_CSPR_ROLE_ARN_IF_SET:
+            # We have either a decorator role or CSPR role, use the effective profile
+            # The generated AWS config will be used here since we set the
+            # AWS_CONFIG_FILE environment variable above.
+            session = boto3.session.Session(profile_name="effective")
         else:
             session = boto3.session.Session(profile_name="task")
     else:
+        # No decorator role or CSPR role, use default session
         # Not using AWS config, just AWS_WEB_IDENTITY_TOKEN_FILE + AWS_ROLE_ARN
         session = boto3.session.Session()
 
-    if role_arn and role_arn != USE_CSPR_ROLE_ARN_IF_SET:
+    if role_arn and role_arn != USE_CSPR_ROLE_ARN_IF_SET and decorator_role_arn != None:
         # If the user provided a role_arn, we assume that role
-        # using the task role credentials. CSPR role is not used.
+        # using the task role credentials. This overrides everything else.
         fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
             client_creator=session._session.create_client,
             source_credentials=session._session.get_credentials(),
@@ -146,8 +155,8 @@ def get_boto3_session(role_arn=None, session_vars=None):
     else:
         # If the user didn't provide a role_arn, or if the role_arn
         # is set to USE_CSPR_ROLE_ARN_IF_SET, we return the default
-        # session which would use the CSPR role if it is set on the
-        # server, and the task role otherwise.
+        # session which would use the effective role (decorator or CSPR) if set,
+        # or the task role otherwise.
         return session
 
 

ob_metaflow_extensions-1.1.163/metaflow_extensions/outerbounds/plugins/aws/__init__.py

@@ -0,0 +1,4 @@
+from .assume_role_decorator import assume_role
+from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
+
+__mf_promote_submodules__ = ["plugins.aws"]

ob_metaflow_extensions-1.1.163/metaflow_extensions/outerbounds/plugins/aws/assume_role.py

@@ -0,0 +1,3 @@
+# Environment variable name used by the assume_role decorator to pass
+# the role ARN to AWS client functions
+OBP_ASSUME_ROLE_ARN_ENV_VAR = "OBP_ASSUME_ROLE_ARN"

ob_metaflow_extensions-1.1.163/metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py

@@ -0,0 +1,65 @@
+from metaflow.user_configs.config_decorators import (
+    MutableFlow,
+    MutableStep,
+    CustomFlowDecorator,
+)
+from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
+
+
+class assume_role(CustomFlowDecorator):
+    """
+    Flow-level decorator for assuming AWS IAM roles.
+
+    When applied to a flow, all steps in the flow will automatically use the specified IAM role-arn
+    as their source principal.
+
+    Usage:
+    ------
+    @assume_role(role_arn="arn:aws:iam::123456789012:role/my-iam-role")
+    class MyFlow(FlowSpec):
+        @step
+        def start(self):
+            import boto3
+            client = boto3.client("dynamodb")  # Automatically uses the role in the flow decorator
+            self.next(self.end)
+
+        @step
+        def end(self):
+            from metaflow import get_aws_client
+            client = get_aws_client("dynamodb")  # Automatically uses the role in the flow decorator
+    """
+
+    def init(self, *args, **kwargs):
+        self.role_arn = kwargs.get("role_arn", None)
+
+        if self.role_arn is None:
+            raise ValueError(
+                "`role_arn` keyword argument is required for the assume_role decorator"
+            )
+
+        if not self.role_arn.startswith("arn:aws:iam::"):
+            raise ValueError(
+                "`role_arn` must be a valid AWS IAM role ARN starting with 'arn:aws:iam::'"
+            )
+
+    def evaluate(self, mutable_flow: MutableFlow) -> None:
+        """
+        This method is called by Metaflow to apply the decorator to the flow.
+        It sets up environment variables that will be used by the AWS client
+        to automatically assume the specified role.
+        """
+        # Import environment decorator at runtime to avoid circular imports
+        from metaflow import environment
+
+        # Set the role ARN as an environment variable that will be picked up
+        # by the get_aws_client function
+        def _setup_role_assumption(step: MutableStep) -> None:
+            # We'll inject the role assumption by adding an environment decorator
+            # The role will be available through an environment variable
+            step.add_decorator(
+                environment, vars={OBP_ASSUME_ROLE_ARN_ENV_VAR: self.role_arn}
+            )
+
+        # Apply the role assumption setup to all steps in the flow
+        for _, step in mutable_flow.steps:
+            _setup_role_assumption(step)

ob_metaflow_extensions-1.1.163/metaflow_extensions/outerbounds/plugins/ollama/__init__.py

@@ -0,0 +1,225 @@
+from metaflow.decorators import StepDecorator
+from metaflow import current
+import functools
+import os
+import threading
+
+from .ollama import OllamaManager, OllamaRequestInterceptor
+from .status_card import OllamaStatusCard
+from ..card_utilities.injector import CardDecoratorInjector
+
+__mf_promote_submodules__ = ["plugins.ollama"]
+
+
+class OllamaDecorator(StepDecorator, CardDecoratorInjector):
+    """
+    This decorator is used to run Ollama APIs as Metaflow task sidecars.
+
+    User code call
+    --------------
+    @ollama(
+        models=[...],
+        ...
+    )
+
+    Valid backend options
+    ---------------------
+    - 'local': Run as a separate process on the local task machine.
+    - (TODO) 'managed': Outerbounds hosts and selects compute provider.
+    - (TODO) 'remote': Spin up separate instance to serve Ollama models.
+
+    Valid model options
+    -------------------
+    Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
+
+    Parameters
+    ----------
+    models: list[str]
+        List of Ollama containers running models in sidecars.
+    backend: str
+        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.
+    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"
+    defaults = {
+        "models": [],
+        "backend": "local",
+        "force_pull": 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(
+        self, flow, graph, step_name, decorators, environment, flow_datastore, logger
+    ):
+        super().step_init(
+            flow, graph, step_name, decorators, environment, flow_datastore, logger
+        )
+        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"],
+                    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:
+                # 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