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

metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_manager.py

@@ -0,0 +1,225 @@
+import os
+import json
+import gzip
+import sys
+import time
+import threading
+import subprocess
+from pathlib import Path
+from typing import Optional, Tuple
+
+import requests
+
+from .constants import (
+    S3_PROXY_BINARY_URLS,
+    DEFAULT_PROXY_PORT,
+    DEFAULT_PROXY_HOST,
+)
+from metaflow.metaflow_config import AWS_SECRETS_MANAGER_DEFAULT_REGION
+from .s3_proxy_api import S3ProxyApiClient
+from .exceptions import S3ProxyException
+
+
+class S3ProxyManager:
+    def __init__(
+        self,
+        integration_name: Optional[str] = None,
+        write_mode: Optional[str] = None,
+        debug: bool = False,
+    ):
+        self.integration_name = integration_name
+        self.write_mode = write_mode
+        self.debug = debug
+        self.process = None
+        self.binary_path = None
+        self.config_path = None
+        self.api_client = S3ProxyApiClient()
+        self.proxy_config = None
+
+    def setup_proxy(self) -> Tuple[dict, int, str, str]:
+        try:
+            if self._is_running_in_kubernetes():
+                config_data = self.api_client.fetch_s3_proxy_config(
+                    self.integration_name
+                )
+                self.binary_path = self._download_binary()
+                self.config_path = self._write_config_file(config_data)
+                # In the new world where the binary is being called
+                # before even the metaflow code exection starts,
+                # so this implies a few important things:
+                # 1, We start the actual proxy process via another python file that safely ships logs to mflog.
+                # 2. We passback the right values to the metaflow step process via env vars.
+                # 3. Metaflow step code relies on env vars to decide if clients need to have s3 proxy in them.
+                self.process = self._start_proxy_process()
+
+                user_code_proxy_config = self._setup_proxy_config(config_data)
+
+                return_tuple = (
+                    user_code_proxy_config,  # this is the config that will be used within the metaflow `step` code.
+                    self.process.pid,  # This is the pid of the process that will jumpstart, monitor and ship logs to MFLOG for the proxy process
+                    self.config_path,  # This is the path to the config that is derived from the integration. It contains the actual bucket path and name where external objects are stored.
+                    self.binary_path,  # This is the path to the binary for the proxy.
+                )
+                # We return a tuple because these values need to be passed down to the metaflow step process where
+                # it will handle thier removal gracefully after the step is finished.
+                return return_tuple
+
+            print(
+                "[@s3_proxy] skipping s3-proxy set up because metaflow has not detected a Kubernetes environment"
+            )
+            raise S3ProxyException(
+                "S3 proxy setup failed because metaflow has not detected a Kubernetes environment"
+            )
+        except Exception as e:
+            if self.debug:
+                print(f"[@s3_proxy] Setup failed: {e}")
+            self.cleanup()
+            raise
+
+    def _is_running_in_kubernetes(self) -> bool:
+        """Check if running inside a Kubernetes pod by checking for Kubernetes service account token."""
+        return (
+            os.path.exists("/var/run/secrets/kubernetes.io/serviceaccount/token")
+            and os.environ.get("KUBERNETES_SERVICE_HOST") is not None
+        )
+
+    def _download_binary(self) -> str:
+        binary_path = Path("/tmp/s3-proxy")
+        if binary_path.exists():
+            if self.debug:
+                print("[@s3_proxy] Binary already exists, skipping download")
+            return str(binary_path.absolute())
+
+        try:
+            if self.debug:
+                print("[@s3_proxy] Downloading binary...")
+
+            from platform import machine
+
+            arch = machine()
+            if arch not in S3_PROXY_BINARY_URLS:
+                raise S3ProxyException(
+                    f"unsupported platform architecture: {arch}. Please reach out to your Outerbounds Support team for more help."
+                )
+
+            response = requests.get(S3_PROXY_BINARY_URLS[arch], stream=True, timeout=60)
+            response.raise_for_status()
+
+            with open(binary_path, "wb") as f:
+                with gzip.GzipFile(fileobj=response.raw) as gz:
+                    f.write(gz.read())
+
+            binary_path.chmod(0o755)
+
+            if self.debug:
+                print("[@s3_proxy] Binary downloaded successfully")
+
+            return str(binary_path.absolute())
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@s3_proxy] Binary download failed: {e}")
+            raise S3ProxyException(f"Failed to download S3 proxy binary: {e}")
+
+    def _write_config_file(self, config_data) -> str:
+        config_path = Path("/tmp/s3-proxy-config.json")
+
+        proxy_config = {
+            "bucketName": config_data.bucket_name,
+            "endpointUrl": config_data.endpoint_url,
+            "accessKeyId": config_data.access_key_id,
+            "accessKeySecret": config_data.secret_access_key,
+            "region": config_data.region,
+        }
+
+        config_path.write_text(json.dumps(proxy_config, indent=2))
+
+        if self.debug:
+            print(f"[@s3_proxy] Config written to {config_path}")
+
+        return str(config_path.absolute())
+
+    def _start_proxy_process(self) -> subprocess.Popen:
+        # This command will jump start a process that will then call the proxy binary
+        # The reason we do something like this is because we need to run all of this before
+        # even the `step` command is called. So we need a python process that will ship the logs
+        # of the proxy process to MFLOG instead of setting print statements. We need this process
+        # to run independently since the S3ProxyManager gets called in the boostrap_proxy which will
+        # exit after jump starting the proxy process.
+        cmd = [self.binary_path, "--bucket-config", self.config_path, "serve"]
+        _env = os.environ.copy()
+        _env["S3_PROXY_BINARY_COMMAND"] = " ".join(cmd)
+        if self.debug:
+            _env["S3_PROXY_BINARY_DEBUG"] = "True"
+        _cmd = [
+            sys.executable,
+            "-m",
+            "metaflow_extensions.outerbounds.plugins.s3_proxy.binary_caller",
+        ]
+        devnull = subprocess.DEVNULL
+        process = subprocess.Popen(
+            _cmd,
+            stdout=devnull,
+            stderr=devnull,
+            text=True,
+            start_new_session=True,
+            env=_env,
+        )
+        time.sleep(3)
+
+        if process.poll() is None:
+            if self.debug:
+                print(f"[@s3_proxy] Proxy started successfully (pid: {process.pid})")
+
+            return process
+        else:
+            stdout_data, stderr_data = process.communicate()
+            if self.debug:
+                print(f"[@s3_proxy] Proxy failed to start - output: {stdout_data}")
+            raise S3ProxyException(f"S3 proxy failed to start: {stdout_data}")
+
+    def _setup_proxy_config(self, config_data):
+        from metaflow.metaflow_config import AWS_SECRETS_MANAGER_DEFAULT_REGION
+
+        region = os.environ.get(
+            "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION",
+            AWS_SECRETS_MANAGER_DEFAULT_REGION,
+        )
+
+        proxy_config = {
+            "endpoint_url": f"http://{DEFAULT_PROXY_HOST}:{DEFAULT_PROXY_PORT}",
+            "region": region,
+            "bucket_name": config_data.bucket_name,
+            "active": True,
+        }
+
+        if self.write_mode:
+            proxy_config["write_mode"] = self.write_mode
+
+        self.proxy_config = proxy_config
+        return proxy_config
+
+    def cleanup(self):
+        try:
+            from metaflow_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import (
+                clear_s3_proxy_config,
+            )
+
+            clear_s3_proxy_config()
+
+            if self.process and self.process.poll() is None:
+                self.process.terminate()
+                self.process.wait(timeout=5)
+                if self.debug:
+                    print("[@s3_proxy] Proxy process stopped")
+
+                from os import remove
+
+                remove(self.config_path)
+                remove(self.binary_path)
+
+        except Exception as e:
+            if self.debug:
+                print(f"[@s3_proxy] Cleanup error: {e}")
+        finally:
+            self.proxy_config = None

metaflow_extensions/outerbounds/plugins/secrets/secrets.py

@@ -0,0 +1,204 @@
+from metaflow.plugins.secrets import SecretsProvider
+from typing import Dict
+
+import base64
+import json
+import requests
+import random
+import time
+import sys
+
+
+class OuterboundsSecretsException(Exception):
+    pass
+
+
+def _api_server_get(*args, conn_error_retries=2, **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 pre-maturely 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 planely 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.
+    """
+    _num_retries = 0
+    noise = random.uniform(-0.5, 0.5)
+    while _num_retries < conn_error_retries:
+        try:
+            return requests.get(*args, **kwargs)
+        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:
+                print(
+                    "[@secrets] Failed to connect to the API server. ",
+                    file=sys.stderr,
+                )
+                raise
+
+
+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 OuterboundsSecretsProvider(SecretsProvider):
+    TYPE = "outerbounds"
+
+    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
+        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 Ourebounds 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:
+            headers = self.headers
+
+        retryable_status_codes = [409]
+        json_payload = json.dumps(payload)
+        for attempt in range(2):  # 0 = initial attempt, 1-2 = retries
+            response = _api_server_get(
+                url, data=json_payload, headers=request_headers, conn_error_retries=5
+            )
+            if response.status_code not in retryable_status_codes:
+                break
+
+            if attempt < 2:  # Don't sleep after the last attempt
+                sleep_time = 0.5 * (attempt + 1)
+                time.sleep(sleep_time)
+
+        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."
+            )
+
+        body = response.json()
+        status_code = body.get("error", {}).get("statusCode", response.status_code)
+        if status_code == 404:
+            raise OuterboundsSecretsException(f"Secret not found: {body}")
+
+        if status_code >= 400:
+            try:
+                raise OuterboundsSecretsException(
+                    f"status_code={status_code}\t*{body['error']['details']['kind']}*\n{body['error']['details']['message']}"
+                )
+            except KeyError:
+                raise OuterboundsSecretsException(
+                    f"status_code={status_code} Unexpected error: {body}"
+                )

metaflow_extensions/outerbounds/plugins/snowflake/__init__.py

@@ -0,0 +1,3 @@
+from .snowflake import connect, get_snowflake_token, Snowflake
+
+__mf_promote_submodules__ = ["plugins.snowflake"]