ob-metaflow-extensions 1.1.45rc3__py2.py3-none-any.whl → 1.5.1__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/apps/core/experimental/__init__.py

@@ -0,0 +1,89 @@
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..app_config import AppConfig
+
+DEFAULT_BRANCH = "test"
+
+
+# Account for project / branch and the capsule input.
+def capsule_input_overrides(app_config: "AppConfig", capsule_input: dict):
+    project = app_config.get_state("project", None)
+    # Update the project/branch related configurations.
+    if project is not None:
+        branch = app_config.get_state("branch", DEFAULT_BRANCH)
+        capsule_input["tags"].extend(
+            [dict(key="project", value=project), dict(key="branch", value=branch)]
+        )
+
+    persistence = app_config.get_state("persistence", None)
+    if persistence is not None and persistence != "none":
+        capsule_input["persistence"] = persistence
+
+    capsule_input["generateStaticUrl"] = app_config.get_state(
+        "generate_static_url", False
+    )
+
+    model_asset_conf = app_config.get_state("models", None)
+    data_asset_conf = app_config.get_state("data", None)
+    code_info = _code_info(app_config)
+    # todo:fix me
+    _objects_key = "associatedObjects"
+    if model_asset_conf or data_asset_conf or code_info:
+        capsule_input[_objects_key] = {}
+
+        if model_asset_conf:
+            capsule_input[_objects_key]["models"] = [
+                {"assetId": x["asset_id"], "assetInstanceId": x["asset_instance_id"]}
+                for x in model_asset_conf
+            ]
+        if data_asset_conf:
+            capsule_input[_objects_key]["data"] = [
+                {"assetId": x["asset_id"], "assetInstanceId": x["asset_instance_id"]}
+                for x in data_asset_conf
+            ]
+        if code_info:
+            capsule_input[_objects_key]["code"] = code_info
+
+    return capsule_input
+
+
+def _code_info(app_config: "AppConfig"):
+    from metaflow.metaflow_git import get_repository_info
+
+    try:
+        from metaflow.metaflow_git import _call_git  # type: ignore
+    except ImportError:
+        # Fallback if _call_git is not available
+        def _call_git(args, path=None):
+            return "", 1, True
+
+    package_dirs = app_config.get_state("packaging_directories")
+    if package_dirs is None:
+        return None
+    for directory in package_dirs:
+        repo_info = get_repository_info(directory)
+        if len(repo_info) == 0:
+            continue
+        git_log_info, returncode, failed = _call_git(
+            ["log", "-1", "--pretty=%B"],
+            path=directory,
+        )
+        repo_url = repo_info["repo_url"]
+        if isinstance(repo_url, str):
+            _url = (
+                repo_url if not repo_url.endswith(".git") else repo_url.rstrip(".git")
+            )
+        else:
+            _url = str(repo_url)
+        _code_info = {
+            "commitId": repo_info["commit_sha"],
+            "commitLink": os.path.join(_url, "commit", str(repo_info["commit_sha"])),
+        }
+        if not failed and returncode == 0 and isinstance(git_log_info, str):
+            _code_info["commitMessage"] = git_log_info.strip()
+
+        return _code_info
+
+    return None

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

@@ -0,0 +1,87 @@
+import os
+import json
+from typing import Tuple, Union
+
+
+class PerimeterExtractor:
+    @classmethod
+    def for_ob_cli(
+        cls, config_dir: str, profile: str
+    ) -> Union[Tuple[str, str], Tuple[None, None]]:
+        """
+        This function will be called when we are trying to extract the perimeter
+        via the ob cli's execution. We will rely on the following logic:
+        1. check environment variables like OB_CURRENT_PERIMETER / OBP_PERIMETER
+        2. run init config to extract the perimeter related configurations.
+
+        Returns
+        -------
+            Tuple[str, str] : Tuple containing perimeter name , API server url.
+        """
+        from outerbounds.utils import metaflowconfig
+
+        perimeter = None
+        api_server = None
+        if os.environ.get("OB_CURRENT_PERIMETER") or os.environ.get("OBP_PERIMETER"):
+            perimeter = os.environ.get("OB_CURRENT_PERIMETER") or os.environ.get(
+                "OBP_PERIMETER"
+            )
+
+        if os.environ.get("OBP_API_SERVER"):
+            api_server = os.environ.get("OBP_API_SERVER")
+
+        if perimeter is None or api_server is None:
+            metaflow_config = metaflowconfig.init_config(config_dir, profile)
+            perimeter = metaflow_config.get("OBP_PERIMETER")
+            api_server = metaflowconfig.get_sanitized_url_from_config(
+                config_dir, profile, "OBP_API_SERVER"
+            )
+
+        return perimeter, api_server  # type: ignore
+
+    @classmethod
+    def during_programmatic_access(cls) -> Union[Tuple[str, str], Tuple[None, None]]:
+        from metaflow.metaflow_config_funcs import init_config
+
+        clean_url = (
+            lambda url: f"https://{url}".rstrip("/")
+            if not url.startswith("https://")
+            else url
+        )
+
+        config = init_config()
+        api_server, perimeter, integrations_url = None, None, None
+        perimeter = config.get(
+            "OBP_PERIMETER", os.environ.get("OBP_PERIMETER", perimeter)
+        )
+        if perimeter is None:
+            raise RuntimeError(
+                "Perimeter not found in metaflow config or environment variables"
+            )
+
+        api_server = config.get(
+            "OBP_API_SERVER", os.environ.get("OBP_API_SERVER", api_server)
+        )
+
+        if api_server is not None and not api_server.startswith("https://"):
+            api_server = clean_url(api_server)
+
+        if api_server is not None:
+            return perimeter, api_server
+
+        integrations_url = config.get(
+            "OBP_INTEGRATIONS_URL", os.environ.get("OBP_INTEGRATIONS_URL", None)
+        )
+
+        if integrations_url is not None and not integrations_url.startswith("https://"):
+            integrations_url = clean_url(integrations_url)
+
+        if integrations_url is not None:
+            api_server = integrations_url.rstrip("/integrations")
+
+        if api_server is None:
+            raise RuntimeError(
+                "API server not found in metaflow config or environment variables"
+            )
+
+        return perimeter, api_server

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

@@ -0,0 +1,164 @@
+from typing import Dict
+
+import base64
+import json
+import requests
+import random
+import time
+import sys
+
+from .utils import safe_requests_wrapper, TODOException
+
+
+class OuterboundsSecretsException(Exception):
+    pass
+
+
+class SecretNotFound(OuterboundsSecretsException):
+    pass
+
+
+class OuterboundsSecretsApiResponse:
+    def __init__(self, response):
+        self.response = response
+
+    @property
+    def secret_resource_id(self):
+        return self.response["secret_resource_id"]
+
+    @property
+    def secret_backend_type(self):
+        return self.response["secret_backend_type"]
+
+
+class SecretRetriever:
+    def get_secret_as_dict(self, secret_id, options={}, role=None):
+        """
+        Supports a special way of specifying secrets sources in outerbounds using the format:
+            @secrets(sources=["outerbounds.<integrations_name>"])
+
+        When invoked it makes a requests to the integrations secrets metadata endpoint on the
+        keywest server to get the cloud resource id for a secret. It then uses that to invoke
+        secrets manager on the core oss and returns the secrets.
+        """
+        headers = {"Content-Type": "application/json", "Connection": "keep-alive"}
+        perimeter, integrations_url = self._get_secret_configs()
+        integration_name = secret_id
+        request_payload = {
+            "perimeter_name": perimeter,
+            "integration_name": integration_name,
+        }
+        response = self._make_request(integrations_url, headers, request_payload)
+        secret_resource_id = response.secret_resource_id
+        secret_backend_type = response.secret_backend_type
+
+        from metaflow.plugins.secrets.secrets_decorator import (
+            get_secrets_backend_provider,
+        )
+
+        secrets_provider = get_secrets_backend_provider(secret_backend_type)
+        secret_dict = secrets_provider.get_secret_as_dict(
+            secret_resource_id, options={}, role=role
+        )
+
+        # Outerbounds stores secrets as binaries. Hence we expect the returned secret to be
+        # {<cloud-secret-name>: <base64 encoded full secret>}. We decode the secret here like:
+        # 1. decode the base64 encoded full secret
+        # 2. load the decoded secret as a json
+        # 3. decode the base64 encoded values in the dict
+        # 4. return the decoded dict
+        binary_secret = next(iter(secret_dict.values()))
+        return self._decode_secret(binary_secret)
+
+    def _is_base64_encoded(self, data):
+        try:
+            if isinstance(data, str):
+                # Check if the string can be base64 decoded
+                base64.b64decode(data).decode("utf-8")
+                return True
+            return False
+        except Exception:
+            return False
+
+    def _decode_secret(self, secret):
+        try:
+            result = {}
+            secret_str = secret
+            if self._is_base64_encoded(secret):
+                # we check if the secret string is base64 encoded because the returned secret from
+                # AWS secret manager is base64 encoded while the secret from GCP is not
+                secret_str = base64.b64decode(secret).decode("utf-8")
+
+            secret_dict = json.loads(secret_str)
+            for key, value in secret_dict.items():
+                result[key] = base64.b64decode(value).decode("utf-8")
+
+            return result
+        except Exception as e:
+            raise OuterboundsSecretsException(f"Error decoding secret: {e}")
+
+    def _get_secret_configs(self):
+        from metaflow_extensions.outerbounds.remote_config import init_config  # type: ignore
+        from os import environ
+
+        conf = init_config()
+        if "OBP_PERIMETER" in conf:
+            perimeter = conf["OBP_PERIMETER"]
+        else:
+            # if the perimeter is not in metaflow config, try to get it from the environment
+            perimeter = environ.get("OBP_PERIMETER", "")
+
+        if "OBP_INTEGRATIONS_URL" in conf:
+            integrations_url = conf["OBP_INTEGRATIONS_URL"]
+        else:
+            # if the integrations is not in metaflow config, try to get it from the environment
+            integrations_url = environ.get("OBP_INTEGRATIONS_URL", "")
+
+        if not perimeter:
+            raise OuterboundsSecretsException(
+                "No perimeter set. Please make sure to run `outerbounds configure <...>` command which can be found on the Outerbounds UI or reach out to your Outerbounds support team."
+            )
+
+        if not integrations_url:
+            raise OuterboundsSecretsException(
+                "No integrations url set. Please notify your Outerbounds support team about this issue."
+            )
+
+        integrations_secrets_metadata_url = f"{integrations_url}/secrets/metadata"
+        return perimeter, integrations_secrets_metadata_url
+
+    def _make_request(self, url, headers: Dict, payload: Dict):
+        try:
+            from metaflow.metaflow_config import SERVICE_HEADERS
+
+            request_headers = {**headers, **(SERVICE_HEADERS or {})}
+        except ImportError:
+            raise OuterboundsSecretsException(
+                "Failed to create app: No Metaflow service headers found"
+            )
+
+        response = safe_requests_wrapper(
+            requests.get,
+            url,
+            data=json.dumps(payload),
+            headers=request_headers,
+            conn_error_retries=5,
+            retryable_status_codes=[409],
+        )
+        self._handle_error_response(response)
+        return OuterboundsSecretsApiResponse(response.json())
+
+    @staticmethod
+    def _handle_error_response(response: requests.Response):
+        if response.status_code >= 500:
+            raise OuterboundsSecretsException(
+                f"Server error: {response.text}. Please reach out to your Outerbounds support team."
+            )
+        status_code = response.status_code
+        if status_code == 404:
+            raise SecretNotFound(f"Secret not found: {response.text}")
+
+        if status_code >= 400:
+            raise OuterboundsSecretsException(
+                f"status_code={status_code}\t\n\t\t{response.text}"
+            )

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

@@ -0,0 +1,233 @@
+import random
+import time
+import sys
+import json
+import requests
+from typing import Optional
+
+# This click import is not used to construct any ob
+# package cli. Its used only for printing stuff.
+# So we can use the static metaflow._vendor import path
+from metaflow._vendor import click
+from .app_config import CAPSULE_DEBUG
+import sys
+import threading
+import time
+import logging
+import itertools
+from typing import Union, Callable, Any, List
+
+from ._vendor.spinner import (
+    Spinners,
+)
+
+
+class MultiStepSpinner:
+    """
+    A spinner that supports multi-step progress and configurable alignment.
+
+    Parameters
+    ----------
+    spinner : Spinners
+        Which spinner frames/interval to use.
+    text : str
+        Static text to display beside the spinner.
+    color : str, optional
+        Click color name.
+    align : {'left','right'}
+        Whether to render the spinner to the left (default) or right of the text.
+    """
+
+    def __init__(
+        self,
+        spinner: Spinners = Spinners.dots,
+        text: str = "",
+        color: Optional[str] = None,
+        align: str = "right",
+        file=sys.stdout,
+    ):
+        cfg = spinner.value
+        self.frames = cfg["frames"]
+        self.interval = float(cfg["interval"]) / 1000.0  # type: ignore
+        self.text = text
+        self.color = color
+        if align not in ("left", "right"):
+            raise ValueError("align must be 'left' or 'right'")
+        self.align = align
+        self._write_file = file
+        # precompute clear length: max frame width + space + text length
+        max_frame = max(self.frames, key=lambda x: len(x))  # type: ignore
+        self.clear_len = len(self.main_text) + len(max_frame) + 1
+
+        self._stop_evt = threading.Event()
+        self._pause_evt = threading.Event()
+        self._thread = None
+        self._write_lock = threading.Lock()
+
+    @property
+    def main_text(self):
+        # if self.text is a callable then call it
+        if callable(self.text):
+            return self.text()
+        return self.text
+
+    def _spin(self):
+        for frame in itertools.cycle(self.frames):
+            if self._stop_evt.is_set():
+                break
+            if self._pause_evt.is_set():
+                time.sleep(0.05)
+                continue
+
+            # ---- Core logging critical section ----
+            with self._write_lock:
+                symbol = click.style(frame, fg=self.color) if self.color else frame
+                if self.align == "left":
+                    msg = f"{symbol} {self.main_text}"
+                else:
+                    msg = f"{self.main_text} {symbol}"
+
+                click.echo(msg, nl=False, file=self._write_file)
+                click.echo("\r", nl=False, file=self._write_file)
+                self._write_file.flush()
+            # ---- End of critical section ----
+            time.sleep(self.interval)
+        # clear the line when done
+        self._clear_line()
+
+    def _clear_line(self):
+        with self._write_lock:
+            click.echo(" " * self.clear_len, nl=False, file=self._write_file)
+            click.echo("\r", nl=False, file=self._write_file)
+            self._write_file.flush()
+
+    def start(self):
+        if self._thread and self._thread.is_alive():
+            return
+        self._stop_evt.clear()
+        self._pause_evt.clear()
+        self._thread = threading.Thread(target=self._spin, daemon=True)
+        self._thread.start()
+
+    def stop(self):
+        self._stop_evt.set()
+        if self._thread:
+            self._thread.join()
+
+    def log(self, *messages: str):
+        """Pause the spinner, emit a ✔ + message, then resume."""
+        self._pause_evt.set()
+        self._clear_line()
+        # ---- Core logging critical section ----
+        with self._write_lock:
+            self._write_file.flush()
+            for message in messages:
+                click.echo(f"{message}", file=self._write_file, nl=True)
+            self._write_file.flush()
+        # ---- End of critical section ----
+        self._pause_evt.clear()
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        self.stop()
+
+
+class SpinnerLogHandler(logging.Handler):
+    def __init__(self, spinner: MultiStepSpinner, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.spinner = spinner
+
+    def emit(self, record):
+        msg = self.format(record)
+        self.spinner.log(msg)
+
+
+class MaximumRetriesExceeded(Exception):
+    def __init__(self, url, method, status_code, text):
+        self.url = url
+        self.method = method
+        self.status_code = status_code
+        self.text = text
+
+    def __str__(self):
+        return f"Maximum retries exceeded for {self.url}[{self.method}] {self.status_code} {self.text}"
+
+
+class TODOException(Exception):
+    pass
+
+
+requests_funcs = [
+    requests.get,
+    requests.post,
+    requests.put,
+    requests.delete,
+    requests.patch,
+    requests.head,
+    requests.options,
+]
+
+
+def safe_requests_wrapper(
+    requests_module_fn,
+    *args,
+    conn_error_retries=2,
+    retryable_status_codes=[409],
+    logger_fn=None,
+    **kwargs,
+):
+    """
+    There are two categories of errors that we need to handle when dealing with any API server.
+    1. HTTP errors. These are are errors that are returned from the API server.
+        - How to handle retries for this case will be application specific.
+    2. Errors when the API server may not be reachable (DNS resolution / network issues)
+        - In this scenario, we know that something external to the API server is going wrong causing the issue.
+        - Failing prematurely in the case might not be the best course of action since critical user jobs might crash on intermittent issues.
+        - So in this case, we can just plainly retry the request.
+
+    This function handles the second case. It's a simple wrapper to handle the retry logic for connection errors.
+    If this function is provided a `conn_error_retries` of 5, then the last retry will have waited 32 seconds.
+    Generally this is a safe enough number of retries after which we can assume that something is really broken. Until then,
+    there can be intermittent issues that would resolve themselves if we retry gracefully.
+    """
+    if requests_module_fn not in requests_funcs:
+        raise ValueError(
+            f"safe_requests_wrapper doesn't support {requests_module_fn.__name__}. You can only use the following functions: {requests_funcs}"
+        )
+
+    _num_retries = 0
+    noise = random.uniform(-0.5, 0.5)
+    response = None
+    while _num_retries < conn_error_retries:
+        try:
+            response = requests_module_fn(*args, **kwargs)
+            if response.status_code not in retryable_status_codes:
+                return response
+            if CAPSULE_DEBUG:
+                if logger_fn:
+                    logger_fn(
+                        f"[outerbounds-debug] safe_requests_wrapper: {response.url}[{requests_module_fn.__name__}] {response.status_code} {response.text}",
+                    )
+                else:
+                    print(
+                        f"[outerbounds-debug] safe_requests_wrapper: {response.url}[{requests_module_fn.__name__}] {response.status_code} {response.text}",
+                        file=sys.stderr,
+                    )
+            _num_retries += 1
+            time.sleep((2 ** (_num_retries + 1)) + noise)
+        except requests.exceptions.ConnectionError:
+            if _num_retries <= conn_error_retries - 1:
+                # Exponential backoff with 2^(_num_retries+1) seconds
+                time.sleep((2 ** (_num_retries + 1)) + noise)
+                _num_retries += 1
+            else:
+                raise
+    raise MaximumRetriesExceeded(
+        response.url,
+        requests_module_fn.__name__,
+        response.status_code,
+        response.text,
+    )

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

@@ -0,0 +1,17 @@
+import os
+from typing import List
+from .app_config import AppConfig, AppConfigError
+from .secrets import SecretRetriever, SecretNotFound
+
+
+def secrets_validator(secrets: List[str]):
+    secret_retriever = SecretRetriever()
+    for secret in secrets:
+        try:
+            secret_retriever.get_secret_as_dict(secret)
+        except SecretNotFound:
+            raise Exception(f"Secret named `{secret}` not found")
+
+
+def run_validations(app_config: AppConfig):
+    pass