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

metaflow_extensions/outerbounds/plugins/apps/core/exceptions.py

@@ -0,0 +1,341 @@
+import json
+from typing import Optional, List, Dict
+from ._state_machine import LogLine
+
+
+class OuterboundsBackendUnhealthyException(Exception):
+    """This exception is raised when the Outerbounds platform is unhealthy (5xx errors) or unreachable."""
+
+    def __init__(
+        self,
+        url: str,
+        method: str,
+        status_code: Optional[int] = None,
+        text: Optional[str] = None,
+        message: Optional[str] = None,
+    ):
+        self.url = url
+        self.method = method
+        self.status_code = status_code
+        self.text = text
+        self.message = message
+        super().__init__(self.message)
+
+
+class OuterboundsForbiddenException(Exception):
+    """This exception is raised when access to an Outerbounds API is forbidden (HTTP 403)."""
+
+    def __init__(
+        self,
+        url: str,
+        method: str,
+        text: str,
+    ):
+        self.url = url
+        self.method = method
+        self.text = text
+        self.message = (
+            f"Access forbidden (HTTP 403) when calling {url}. "
+            "This typically means your credentials lack permission for this operation. "
+            "Please verify that you outerbounds / metaflow configuration has access to the "
+            "correct perimeter.\n"
+            "If you believe this is an error, contact your Outerbounds administrator."
+        )
+        super().__init__(self.message)
+
+
+class OuterboundsConfigurationException(Exception):
+    """This exception is raised when Outerbounds configuration is missing or invalid."""
+
+    def __init__(self, missing_config: str):
+        self.missing_config = missing_config
+        self.message = (
+            f"Outerbounds configuration '{missing_config}' not found.\n\n"
+            "If running locally:\n"
+            "  - Run: outerbounds configure <magic-string-from-outerbounds-ui>\n"
+            "If running remotely (e.g., in a Metaflow task):\n"
+            "  - Ensure you have the Outerbounds distribution installed (pip install outerbounds)\n"
+            "    and not the open-source Metaflow. The Outerbounds fork injects configuration\n"
+            "    into remote tasks automatically which may not be present in open source metaflow.\n"
+        )
+        super().__init__(self.message)
+
+
+class CapsuleApiException(Exception):
+    def __init__(
+        self,
+        url: str,
+        method: str,
+        status_code: int,
+        text: str,
+        message: Optional[str] = None,
+    ):
+        self.url = url
+        self.method = method
+        self.status_code = status_code
+        self.text = text
+        self.message = message
+
+    def __str__(self):
+        return (
+            f"CapsuleApiException: {self.url} [{self.method}]: Status Code: {self.status_code} \n\n {self.text}"
+            + (f"\n\n {self.message}" if self.message else "")
+        )
+
+
+class CapsuleDeploymentException(Exception):
+    """Base exception for all capsule deployment failures."""
+
+    def __init__(
+        self,
+        capsule_id: str,
+        message: str,
+    ):
+        self.capsule_id = capsule_id
+        self.message = message
+        super().__init__(self.message)
+
+    def __str__(self):
+        return f"CapsuleDeploymentException: [{self.capsule_id}] :: {self.message}"
+
+
+class CapsuleCrashLoopException(CapsuleDeploymentException):
+    """Raised when a worker enters CrashLoopBackOff or Failed state."""
+
+    def __init__(
+        self,
+        capsule_id: str,
+        worker_id: str,
+        logs: Optional[List[LogLine]] = None,
+    ):
+        self.worker_id = worker_id
+        self.logs = logs or []
+        message = f"Worker ID ({worker_id}) is crashlooping. Please check the logs for more information."
+        super().__init__(capsule_id, message)
+
+    def __str__(self):
+        return f"CapsuleCrashLoopException: [{self.capsule_id}] :: {self.message}"
+
+
+class CapsuleReadinessException(CapsuleDeploymentException):
+    """Raised when capsule fails to meet readiness conditions within timeout."""
+
+    def __init__(
+        self,
+        capsule_id: str,
+        reason: Optional[str] = None,
+    ):
+        message = f"Capsule {capsule_id} failed to be ready to serve traffic"
+        if reason:
+            message += f": {reason}"
+        super().__init__(capsule_id, message)
+
+    def __str__(self):
+        return f"CapsuleReadinessException: [{self.capsule_id}] :: {self.message}"
+
+
+class CapsuleConcurrentUpgradeException(CapsuleDeploymentException):
+    """Raised when a concurrent upgrade invalidates the current deployment."""
+
+    def __init__(
+        self,
+        capsule_id: str,
+        expected_version: str,
+        actual_version: str,
+        modified_by: Optional[str] = None,
+        modified_at: Optional[str] = None,
+    ):
+        self.expected_version = expected_version
+        self.actual_version = actual_version
+        self.modified_by = modified_by
+        self.modified_at = modified_at
+        message = (
+            f"A capsule upgrade was triggered outside current deployment instance. "
+            f"Expected version: {expected_version}, actual version: {actual_version}"
+        )
+        if modified_by:
+            message += f". Modified by: {modified_by}"
+        if modified_at:
+            message += f" at {modified_at}"
+        super().__init__(capsule_id, message)
+
+    def __str__(self):
+        return (
+            f"CapsuleConcurrentUpgradeException: [{self.capsule_id}] :: {self.message}"
+        )
+
+
+class CapsuleDeletedDuringDeploymentException(CapsuleDeploymentException):
+    """Raised when a capsule is deleted while deployment is in progress."""
+
+    def __init__(self, capsule_id: str):
+        super().__init__(capsule_id, "Capsule was deleted during deployment")
+
+
+class CodePackagingException(Exception):
+    """Exception raised when code packaging fails."""
+
+    pass
+
+
+class AppNotFoundException(Exception):
+    pass
+
+
+class AppCreationFailedException(Exception):
+    """Raised when app deployment submission fails due to an API error."""
+
+    def __init__(
+        self,
+        app_name: str,
+        status_code: int,
+        error_text: str,
+    ):
+        self.status_code = status_code
+        self.error_text = error_text
+        message = f"Failed to submit app deployment: HTTP {status_code} - {error_text}"
+        if status_code == 400:
+            # 400 = validation error; the app configuration is invalid and must be fixed.
+            message = "Invalid app deployment configuration submitted: "
+            try:
+                reason_for_failure = json.loads(error_text).get("message", error_text)
+            except json.JSONDecodeError:
+                reason_for_failure = error_text
+            message += reason_for_failure
+            message += (
+                "\n\nCheck your config file or CLI parameters if deploying via CLI, "
+                "or AppDeployer parameters if deploying programmatically."
+            )
+
+        self.message = message
+        super().__init__(message)
+
+    def __str__(self):
+        return f"AppCreationFailedException: {self.message}"
+
+
+class AppDeploymentException(Exception):
+    """Base exception for all individual app deployment failures."""
+
+    def __init__(self, app_id: str, message: str):
+        self.app_id = app_id
+        self.message = message
+        self._deployed_app = None
+        super().__init__(self.message)
+
+    def __str__(self):
+        return f"AppDeploymentException: [{self.app_id}] :: {self.message}"
+
+    @property
+    def deployed_app(self):
+        from .deployer import DeployedApp
+
+        return DeployedApp._from_capsule_id(self.app_id)
+
+
+class AppCrashLoopException(AppDeploymentException):
+    """Raised when an app worker enters CrashLoopBackOff or Failed state."""
+
+    def __init__(
+        self,
+        app_id: str,
+        worker_id: str,
+        logs: Optional[List] = None,
+    ):
+        self.worker_id = worker_id
+        self.logs = logs or []
+        message = f"Worker ({worker_id}) is crashlooping. Please check the logs for more information."
+        super().__init__(app_id, message)
+
+    def __str__(self):
+        return f"AppCrashLoopException: [{self.app_id}] :: {self.message}"
+
+
+class AppReadinessException(AppDeploymentException):
+    """Raised when app fails to meet readiness conditions within timeout."""
+
+    def __init__(self, app_id: str, reason: Optional[str] = None):
+        message = f"App {app_id} failed to be ready to serve traffic"
+        if reason:
+            message += f": {reason}"
+        super().__init__(app_id, message)
+
+    def __str__(self):
+        return f"AppReadinessException: [{self.app_id}] :: {self.message}"
+
+
+class AppUpgradeInProgressException(AppDeploymentException):
+    """Raised when attempting to deploy while another upgrade is already in progress."""
+
+    def __init__(
+        self,
+        app_id: str,
+        upgrader: Optional[str] = None,
+    ):
+        self.upgrader = upgrader
+        if upgrader:
+            message = (
+                f"App {app_id} is currently being upgraded by {upgrader}. "
+                "Use force_upgrade=True in AppDeployer to override."
+            )
+        else:
+            message = (
+                f"App {app_id} is currently being upgraded. "
+                "Use force_upgrade=True in AppDeployer to override."
+            )
+        super().__init__(app_id, message)
+
+    def __str__(self):
+        return f"AppUpgradeInProgressException: [{self.app_id}] :: {self.message}"
+
+
+class AppConcurrentUpgradeException(AppDeploymentException):
+    """Raised when a concurrent upgrade invalidates the current deployment mid-flight."""
+
+    def __init__(
+        self,
+        app_id: str,
+        expected_version: str,
+        actual_version: str,
+        modified_by: Optional[str] = None,
+        modified_at: Optional[str] = None,
+    ):
+        self.expected_version = expected_version
+        self.actual_version = actual_version
+        self.modified_by = modified_by
+        self.modified_at = modified_at
+
+        modifier_info = ""
+        if modified_by:
+            modifier_info = f" by '{modified_by}'"
+            if modified_at:
+                modifier_info += f" at {modified_at}"
+
+        message = (
+            f"Another deployment was triggered{modifier_info} while this deployment was in progress.\n\n"
+            f"This deployment expected to be working with version '{expected_version}', but the app "
+            f"is now at version '{actual_version}'. The current deployment has been invalidated.\n\n"
+            "To avoid this in the future, you can either use a unique `name` for each deployment "
+            "or coordinate deployments to ensure concurrent upgrades to the same app don't overlap."
+        )
+        super().__init__(app_id, message)
+
+    def __str__(self):
+        return f"AppConcurrentUpgradeException: [{self.app_id}] :: {self.message}"
+
+
+class AppDeletedDuringDeploymentException(AppDeploymentException):
+    """Raised when an app is deleted while deployment is in progress."""
+
+    def __init__(self, app_id: str):
+        message = (
+            f"App '{app_id}' was deleted while this deployment was in progress.\n\n"
+            "This can happen when another process or user deletes the app during deployment. "
+            "Since apps can be programmatically created and deleted, concurrent operations "
+            "may conflict. If you did not intend to delete this app, check for other processes "
+            "or users that may be managing deployments in this perimeter."
+        )
+        super().__init__(app_id, message)
+
+    def __str__(self):
+        return f"AppDeletedDuringDeploymentException: [{self.app_id}] :: {self.message}"

metaflow_extensions/outerbounds/plugins/apps/core/perimeters.py

@@ -2,8 +2,15 @@ import os
 import json
 from typing import Tuple, Union
 
+import requests
+from .utils import safe_requests_wrapper
+from .exceptions import OuterboundsConfigurationException
+
 
 class PerimeterExtractor:
+
+    config = None
+
     @classmethod
     def for_ob_cli(
         cls, config_dir: str, profile: str
@@ -40,7 +47,7 @@ class PerimeterExtractor:
         return perimeter, api_server  # type: ignore
 
     @classmethod
-    def during_metaflow_execution(cls) -> Union[Tuple[str, str], Tuple[None, None]]:
+    def during_programmatic_access(cls) -> Union[Tuple[str, str], Tuple[None, None]]:
         from metaflow.metaflow_config_funcs import init_config
 
         clean_url = (
@@ -55,9 +62,7 @@ class PerimeterExtractor:
             "OBP_PERIMETER", os.environ.get("OBP_PERIMETER", perimeter)
         )
         if perimeter is None:
-            raise RuntimeError(
-                "Perimeter not found in metaflow config or environment variables"
-            )
+            raise OuterboundsConfigurationException("OBP_PERIMETER")
 
         api_server = config.get(
             "OBP_API_SERVER", os.environ.get("OBP_API_SERVER", api_server)
@@ -80,8 +85,39 @@ class PerimeterExtractor:
             api_server = integrations_url.rstrip("/integrations")
 
         if api_server is None:
+            raise OuterboundsConfigurationException("OBP_API_SERVER")
+
+        return perimeter, api_server
+
+    @classmethod
+    def config_during_programmatic_access(cls) -> dict:
+        #!HACK: Resolving remote configs is a PITA (all the variable piping we need to do via metaflow)
+        # So instead we will just derive the URL. We are in this situation because its a pain
+        # to load configurations at arbitrary points in the runtime.
+        if cls.config is not None:
+            return json.loads(json.dumps(cls.config))  # Return fresh copy
+        from metaflow.metaflow_config import SERVICE_HEADERS
+
+        perimeter, api_server = cls.during_programmatic_access()
+        response = safe_requests_wrapper(
+            requests.get,
+            f"{api_server}/v1/perimeters/{perimeter}/metaflowconfigs/default",
+            headers=SERVICE_HEADERS,
+        )
+        if response.status_code >= 400:
             raise RuntimeError(
-                "API server not found in metaflow config or environment variables"
+                f"Server error: {response.text}. Please reach out to your Outerbounds support team."
             )
+        try:
+            remote_config = response.json()
 
-        return perimeter, api_server
+            if not remote_config.get("config"):
+                raise json.JSONDecodeError
+        except json.JSONDecodeError:
+            raise RuntimeError(
+                "Exception retrieving remote outerbounds configuration. "
+                "Please reach out to Outerbounds suport team with this stack trace."
+            )
+
+        cls.config = remote_config.get("config")
+        return json.loads(json.dumps(cls.config))  # Return fresh copy

metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py

@@ -25,10 +25,29 @@ class assume_role(FlowMutator):
         def end(self):
             from metaflow import get_aws_client
             client = get_aws_client("dynamodb")  # Automatically uses the role in the flow decorator
+
+    You can also filter which steps should use the role:
+    @assume_role(role_arn="arn:aws:iam::123456789012:role/my-iam-role", steps=["start", "process"])
+    class MyFlow(FlowSpec):
+        @step
+        def start(self):
+            # user code in this step will use the assumed role
+            pass
+
+        @step
+        def process(self):
+            # user code in this step will use the assumed role
+            pass
+
+        @step
+        def end(self):
+            # user code in this step will NOT use the assumed role
+            pass
     """
 
     def init(self, *args, **kwargs):
         self.role_arn = kwargs.get("role_arn", None)
+        self.steps = kwargs.get("steps", None)
 
         if self.role_arn is None:
             raise ValueError(
@@ -40,6 +59,13 @@ class assume_role(FlowMutator):
                 "`role_arn` must be a valid AWS IAM role ARN starting with 'arn:aws:iam::'"
             )
 
+        # Validate steps parameter
+        if self.steps is not None:
+            if not isinstance(self.steps, (list, tuple)):
+                raise ValueError("`steps` must be a list or tuple of step names")
+            if not all(isinstance(s, str) for s in self.steps):
+                raise ValueError("All step names in `steps` must be strings")
+
     def pre_mutate(self, mutable_flow: MutableFlow) -> None:
         """
         This method is called by Metaflow to apply the decorator to the flow.
@@ -49,6 +75,18 @@ class assume_role(FlowMutator):
         # Import environment decorator at runtime to avoid circular imports
         from metaflow import environment
 
+        # Validate that all specified steps exist in the flow
+        if self.steps is not None:
+            flow_step_names = {step_name for step_name, _ in mutable_flow.steps}
+            specified_steps = set(self.steps)
+            missing_steps = specified_steps - flow_step_names
+
+            if missing_steps:
+                raise ValueError(
+                    f"Step(s) {sorted(missing_steps)} specified in `steps` parameter "
+                    f"do not exist in the flow. Available steps: {sorted(flow_step_names)}"
+                )
+
         def _swap_environment_variables(step: MutableStep, role_arn: str) -> None:
             _step_has_env_set = True
             _env_kwargs = {OBP_ASSUME_ROLE_ARN_ENV_VAR: role_arn}
@@ -73,6 +111,8 @@ class assume_role(FlowMutator):
         def _setup_role_assumption(step: MutableStep) -> None:
             _swap_environment_variables(step, self.role_arn)
 
-        # Apply the role assumption setup to all steps in the flow
-        for _, step in mutable_flow.steps:
-            _setup_role_assumption(step)
+        # Apply the role assumption setup to all steps in the flow (or filtered steps)
+        for step_name, step in mutable_flow.steps:
+            # If steps filter is specified, only apply to those steps
+            if self.steps is None or step_name in self.steps:
+                _setup_role_assumption(step)

metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -27,6 +27,7 @@ def bake_image(
     conda_packages: Optional[Dict[str, str]] = None,
     base_image: Optional[str] = None,
     logger: Optional[Callable[[str], Any]] = None,
+    fast_bakery_url: Optional[str] = None,
 ) -> FastBakeryApiResponse:
     """
     Bakes a Docker image with the specified dependencies.
@@ -39,6 +40,7 @@ def bake_image(
         conda_packages: Dictionary of Conda packages and versions
         base_image: Base Docker image to use
         logger: Optional logger function to output progress
+        fast_bakery_url: Optional FB URL
 
     Returns:
         FastBakeryApiResponse: The response from the bakery service
@@ -50,6 +52,13 @@ def bake_image(
     if logger is None:
         logger = partial(print, file=sys.stderr)
 
+    if all([fast_bakery_url is None and FAST_BAKERY_URL is None]):
+        raise BakerException(
+            "Image Bakery endpoint missing. METAFLOW_FAST_BAKERY_URL environment/configuration variable not found."
+        )
+
+    fast_bakery_url = fast_bakery_url or FAST_BAKERY_URL
+
     # Thread lock for logging
     logger_lock = threading.Lock()
     images_baked = 0
@@ -63,7 +72,7 @@ def bake_image(
         base_image=None,
     ):
         try:
-            bakery = FastBakery(url=FAST_BAKERY_URL)
+            bakery = FastBakery(url=fast_bakery_url)
             bakery._reset_payload()
             bakery.python_version(python)
             bakery.pypi_packages(pypi_packages)

metaflow_extensions/outerbounds/plugins/optuna/__init__.py

@@ -19,7 +19,8 @@ def get_deployment_db_access_endpoint(name: str):
     from ..apps.core.perimeters import PerimeterExtractor
     from ..apps.core.capsule import CapsuleApi
 
-    perimeter, cap_url = PerimeterExtractor.during_metaflow_execution()
+    perimeter, cap_url = PerimeterExtractor.during_programmatic_access()
+    # TODO: Find more natural way to get rid of CapsuleApi here.
     deployment = CapsuleApi(cap_url, perimeter).get_by_name(name)
     if not deployment:
         raise Exception(f"No app deployment found with name `{name}`")

metaflow_extensions/outerbounds/plugins/snowflake/snowflake.py

@@ -83,22 +83,32 @@ def get_snowflake_token(user: str = "", role: str = "", integration: str = "") -
     return response.json()["token"]
 
 
-def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
+def get_oauth_connection_params(
+    user: str = "", role: str = "", integration: str = "", **kwargs
+) -> Dict:
     """
-    Connect to snowflake using the token minted by Outerbounds
+    Get OAuth connection parameters for Snowflake authentication using Outerbounds integration.
+
+    This is a helper function that returns connection parameters dict that can be used
+    with both snowflake-connector-python and snowflake-snowpark-python.
+
     user: str
         The user name used to authenticate with snowflake
     role: str
-        The role to request when connect with snowflake
+        The role to request when connecting with snowflake
     integration: str
-        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists in the current perimeter. If integration is not set and more than one exists in the current perimeter, then we raise an exception.
+        The name of the snowflake integration to use. If not set, an existing integration
+        will be used provided that only one exists in the current perimeter.
     kwargs: dict
-        Additional arguments to pass to the python snowflake connector
+        Additional arguments to include in the connection parameters
+
+    Returns:
+        Dict with connection parameters including OAuth token
     """
     # ensure password is not set
     if "password" in kwargs:
         raise OuterboundsSnowflakeConnectorException(
-            "Password should not be set when using Outerbounds snowflake connector."
+            "Password should not be set when using Outerbounds OAuth authentication."
         )
 
     provisioner = SnowflakeIntegrationProvisioner(integration)
@@ -137,11 +147,31 @@ def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
     kwargs["role"] = role
     kwargs["user"] = user
 
+    return kwargs
+
+
+def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
+    """
+    Connect to snowflake using the token minted by Outerbounds
+    user: str
+        The user name used to authenticate with snowflake
+    role: str
+        The role to request when connect with snowflake
+    integration: str
+        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists in the current perimeter. If integration is not set and more than one exists in the current perimeter, then we raise an exception.
+    kwargs: dict
+        Additional arguments to pass to the python snowflake connector
+    """
+    # Get OAuth connection params using the helper
+    connection_params = get_oauth_connection_params(
+        user=user, role=role, integration=integration, **kwargs
+    )
+
     # connect to snowflake
     try:
         from snowflake.connector import connect
 
-        cn = connect(**kwargs)
+        cn = connect(**connection_params)
         return cn
     except ImportError as ie:
         raise OuterboundsSnowflakeConnectorException(

metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py

@@ -1,9 +1,11 @@
 import os
+import re
 import shlex
 import atexit
 import json
 import math
 import time
+import hashlib
 
 from metaflow import util
 
@@ -57,21 +59,29 @@ class Snowpark(object):
         atexit.register(lambda: self.job.kill() if hasattr(self, "job") else None)
 
     def _job_name(self, user, flow_name, run_id, step_name, task_id, retry_count):
-        return "{user}-{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}".format(
-            user=user,
-            flow_name=flow_name,
-            run_id=str(run_id) if run_id is not None else "",
-            step_name=step_name,
-            task_id=str(task_id) if task_id is not None else "",
-            retry_count=str(retry_count) if retry_count is not None else "",
+        unique_str = (
+            "{user}-{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}".format(
+                user=user,
+                flow_name=flow_name,
+                run_id=str(run_id) if run_id is not None else "",
+                step_name=step_name,
+                task_id=str(task_id) if task_id is not None else "",
+                retry_count=str(retry_count) if retry_count is not None else "",
+            )
         )
+        unique_hash = hashlib.md5(unique_str.encode("utf-8")).hexdigest()[:8]
+        raw_prefix = f"{flow_name}-{step_name}"
+        safe_prefix = re.sub(r"[^a-z0-9]", "-", raw_prefix.lower())
+        safe_prefix = safe_prefix[:54]
+        safe_prefix = safe_prefix.lstrip("-")
+        return f"{safe_prefix}-{unique_hash}"
 
     def _command(self, environment, code_package_url, step_name, step_cmds, task_spec):
         mflog_expr = export_mflog_env_vars(
             datastore_type=self.datastore.TYPE,
             stdout_path=STDOUT_PATH,
             stderr_path=STDERR_PATH,
-            **task_spec
+            **task_spec,
         )
         init_cmds = environment.get_package_commands(
             code_package_url, self.datastore.TYPE