ob-metaflow-extensions 1.1.170__py2.py3-none-any.whl → 1.4.35__py2.py3-none-any.whl

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

@@ -0,0 +1,303 @@
+from .config import TypedCoreConfig, TypedDict
+from .perimeters import PerimeterExtractor
+from .capsule import CapsuleApi
+import json
+from ._state_machine import DEPLOYMENT_READY_CONDITIONS, LogLine
+from .app_config import AppConfig, AppConfigError
+from .capsule import CapsuleDeployer, list_and_filter_capsules
+from functools import partial
+import sys
+import uuid
+from typing import Type, Dict, List
+from datetime import datetime
+
+
+class AppDeployer(TypedCoreConfig):
+    """ """
+
+    __init__ = TypedCoreConfig.__init__
+
+    _app_config: AppConfig
+
+    _state = {}
+
+    __state_items = [
+        "perimeter",
+        "api_url",
+        "code_package_url",
+        "code_package_key",
+        "image",
+        "project",
+        "branch",
+    ]
+
+    @property
+    def _deploy_config(self) -> AppConfig:
+        if not hasattr(self, "_app_config"):
+            self._app_config = AppConfig(self._config)
+        return self._app_config
+
+    # Things that need to be set before deploy
+    @classmethod
+    def _set_state(
+        cls,
+        perimeter: str,
+        api_url: str,
+        code_package_url: str = None,
+        code_package_key: str = None,
+        name_prefix: str = None,
+        image: str = None,
+        max_entropy: int = 4,
+        default_tags: List[Dict[str, str]] = None,
+        project: str = None,
+        branch: str = None,
+    ):
+        cls._state["perimeter"] = perimeter
+        cls._state["api_url"] = api_url
+        cls._state["code_package_url"] = code_package_url
+        cls._state["code_package_key"] = code_package_key
+        cls._state["name_prefix"] = name_prefix
+        cls._state["image"] = image
+        cls._state["max_entropy"] = max_entropy
+        cls._state["default_tags"] = default_tags
+        cls._state["project"] = project
+        cls._state["branch"] = branch
+
+        assert (
+            max_entropy > 0
+        ), "max_entropy must be greater than 0. Since AppDeployer's deploy fn can be called many time inside a step itself."
+
+    def deploy(
+        self,
+        readiness_condition=DEPLOYMENT_READY_CONDITIONS.ATLEAST_ONE_RUNNING,
+        max_wait_time=600,
+        readiness_wait_time=10,
+        logger_fn=partial(print, file=sys.stderr),
+        status_file=None,
+        no_loader=False,
+        **kwargs,
+    ) -> "DeployedApp":
+
+        # Name setting from top level if none is set in the code
+        if self._deploy_config._core_config.name is None:
+            name = self._state[
+                "name_prefix"
+            ]  # for now the name-prefix cannot be very large.
+            entropy = uuid.uuid4().hex[: self._state["max_entropy"]]
+            self._deploy_config._core_config.name = f"{name}-{entropy}"
+
+        if len(self._state["default_tags"]) > 0:
+            self._deploy_config._core_config.tags = (
+                self._deploy_config._core_config.tags or []
+            ) + self._state["default_tags"]
+
+        self._deploy_config.commit()
+        # Set any state that might have been passed down from the top level
+        for k in self.__state_items:
+            if self._deploy_config.get_state(k) is None:
+                self._deploy_config.set_state(k, self._state[k])
+
+        capsule = CapsuleDeployer(
+            self._deploy_config,
+            self._state["api_url"],
+            create_timeout=max_wait_time,
+            debug_dir=None,
+            success_terminal_state_condition=readiness_condition,
+            readiness_wait_time=readiness_wait_time,
+            logger_fn=logger_fn,
+        )
+
+        currently_present_capsules = list_and_filter_capsules(
+            capsule.capsule_api,
+            None,
+            None,
+            capsule.name,
+            None,
+            None,
+            None,
+        )
+
+        force_upgrade = self._deploy_config.get_state("force_upgrade", False)
+
+        if len(currently_present_capsules) > 0:
+            # Only update the capsule if there is no upgrade in progress
+            # Only update a "already updating" capsule if the `--force-upgrade` flag is provided.
+            _curr_cap = currently_present_capsules[0]
+            this_capsule_is_being_updated = _curr_cap.get("status", {}).get(
+                "updateInProgress", False
+            )
+
+            if this_capsule_is_being_updated and not force_upgrade:
+                _upgrader = _curr_cap.get("metadata", {}).get("lastModifiedBy", None)
+                message = f"{capsule.capsule_type} is currently being upgraded"
+                if _upgrader:
+                    message = (
+                        f"{capsule.capsule_type} is currently being upgraded. Upgrade was launched by {_upgrader}. "
+                        "If you wish to force upgrade, you can do so by providing the `--force-upgrade` flag."
+                    )
+                raise AppConfigError(message)
+
+            logger_fn(
+                f"🚀 {'' if not force_upgrade else 'Force'} Upgrading {capsule.capsule_type.lower()} `{capsule.name}`....",
+            )
+        else:
+            logger_fn(
+                f"🚀 Deploying {capsule.capsule_type.lower()} `{capsule.name}`....",
+            )
+
+        capsule.create()
+        final_status = capsule.wait_for_terminal_state()
+        return DeployedApp(
+            final_status["id"],
+            final_status["auth_type"],
+            final_status["public_url"],
+            final_status["name"],
+            final_status["deployed_version"],
+            final_status["deployed_at"],
+        )
+
+
+class DeployedApp:
+    def __init__(
+        self,
+        _id: str,
+        capsule_type: str,
+        public_url: str,
+        name: str,
+        deployed_version: str,
+        deployed_at: str,
+    ):
+        self._id = _id
+        self._capsule_type = capsule_type
+        self._public_url = public_url
+        self._name = name
+        self._deployed_version = deployed_version
+        self._deployed_at = deployed_at
+
+    def _get_capsule_api(self) -> CapsuleApi:
+        perimeter, api_server = PerimeterExtractor.during_metaflow_execution()
+        return CapsuleApi(api_server, perimeter)
+
+    def logs(self, previous=False) -> Dict[str, List[LogLine]]:
+        """
+        Returns a dictionary of worker_id to logs.
+        If `previous` is True, it will return the logs from the previous execution of the workers. Useful when debugging a crashlooping worker.
+        """
+        capsule_api = self._get_capsule_api()
+        # extract workers from capsule
+        workers = capsule_api.get_workers(self._id)
+        # get logs from workers
+        logs = {
+            # worker_id: logs
+        }
+        for worker in workers:
+            # TODO: Handle exceptions better over here.
+            logs[worker["workerId"]] = capsule_api.logs(
+                self._id, worker["workerId"], previous=previous
+            )
+        return logs
+
+    def info(self) -> dict:
+        """
+        Returns a dictionary representing the deployed app.
+        """
+        capsule_api = self._get_capsule_api()
+        capsule = capsule_api.get(self._id)
+        return capsule
+
+    def replicas(self):
+        capsule_api = self._get_capsule_api()
+        return capsule_api.get_workers(self._id)
+
+    def scale_to_zero(self):
+        """
+        Scales the DeployedApp to 0 replicas.
+        """
+        capsule_api = self._get_capsule_api()
+        return capsule_api.patch(
+            self._id,
+            {
+                "autoscalingConfig": {
+                    "minReplicas": 0,
+                    "maxReplicas": 0,
+                }
+            },
+        )
+
+    def delete(self):
+        """
+        Deletes the DeployedApp.
+        """
+        capsule_api = self._get_capsule_api()
+        return capsule_api.delete(self._id)
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def auth_style(self) -> str:
+        # TODO : Fix naming here.
+        return self._capsule_type
+
+    @property
+    def public_url(self) -> str:
+        return self._public_url
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def deployed_version(self) -> str:
+        return self._deployed_version
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self._id,
+            "auth_style": self.auth_style,  # TODO : Fix naming here.
+            "public_url": self._public_url,
+            "name": self._name,
+            "deployed_version": self._deployed_version,
+            "deployed_at": self._deployed_at,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(
+            _id=data["id"],
+            capsule_type=data["capsule_type"],
+            public_url=data["public_url"],
+            name=data["name"],
+            deployed_version=data["deployed_version"],
+            deployed_at=data["deployed_at"],
+        )
+
+    @property
+    def deployed_at(self) -> datetime:
+        return datetime.fromisoformat(self._deployed_at)
+
+    def __repr__(self) -> str:
+        return (
+            f"DeployedApp(id='{self._id}', "
+            f"name='{self._name}', "
+            f"public_url='{self._public_url}', "
+            f"deployed_version='{self._deployed_version}')"
+        )
+
+
+class apps:
+
+    _name_prefix = None
+
+    @classmethod
+    def set_name_prefix(cls, name_prefix: str):
+        cls._name_prefix = name_prefix
+
+    @property
+    def name_prefix(self) -> str:
+        return self._name_prefix
+
+    @property
+    def Deployer(self) -> Type[AppDeployer]:
+        return AppDeployer

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_metaflow_execution(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}"
+            )