ob-metaflow-extensions 1.3.0__tar.gz → 1.3.1__tar.gz

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob_metaflow_extensions
-Version: 1.3.0
+Version: 1.3.1
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/__init__.py

@@ -340,6 +340,9 @@ STEP_DECORATORS_DESC = [
     ("ollama", ".ollama.OllamaDecorator"),
     ("vllm", ".vllm.VLLMDecorator"),
     ("app_deploy", ".apps.app_deploy_decorator.AppDeployDecorator"),
+    ("s3_proxy", ".s3_proxy.s3_proxy_decorator.S3ProxyDecorator"),
+    ("nebius_s3_proxy", ".s3_proxy.s3_proxy_decorator.NebiusS3ProxyDecorator"),
+    ("coreweave_s3_proxy", ".s3_proxy.s3_proxy_decorator.CoreWeaveS3ProxyDecorator"),
 ]
 
 TOGGLE_STEP_DECORATOR = [

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/_state_machine.py

@@ -199,29 +199,29 @@ class DEPLOYMENT_READY_CONDITIONS:
     This allows users or platform designers to configure the criteria for deployment readiness.
 
     Why do we need deployment readiness conditions?
-    - Deployments might be taking place from a CI/CD esq environment, In these setups, the downstream build triggers might be depending on a specific criteria for deployment completion. Having readiness conditions allows the CI/CD systems to get a signal of when the deployment is ready.
+        - Deployments might be taking place from a CI/CD-esque environment, In these setups, the downstream build triggers might be depending on a specific criteria for deployment completion. Having readiness conditions allows the CI/CD systems to get a signal of when the deployment is ready.
     - Users might be calling the deployment API under different conditions:
         - Some users might want a cluster of workers ready before serving traffic while others might want just one worker ready to start serving traffic.
 
     Some readiness conditions include:
-    1) [at_least_one_running] Atleast min(min_replicas, 1) workers of the current deployment instance's version have started running.
+            1) [at_least_one_running] At least min(min_replicas, 1) workers of the current deployment instance's version have started running.
         - Usecase: Some endpoints may be deployed ephemerally and are considered ready when at least one instance is running; additional instances are for load management.
-    2) [all_running] Atleast min_replicas number of workers are running for the deployment to be considered ready.
+        2) [all_running] At least min_replicas number of workers are running for the deployment to be considered ready.
         - Usecase: Operators may require that all replicas are available before traffic is routed. Needed when inference endpoints maybe under some SLA or require a larger load
-    3) [fully_finished] Atleast min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
+        3) [fully_finished] At least min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
         - Usecase: Ensuring endpoint is fully available and no other versions are running or endpoint has been fully scaled down.
     4) [async] The deployment will be assumed ready as soon as the server responds with a 200.
         - Usecase: Operators may only care that the URL is minted for the deployment or the deployment eventually scales down to 0.
     """
 
-    # `ATLEAST_ONE_RUNNING` implies that atleast one worker of the current deployment instance's version has started running.
+    # `ATLEAST_ONE_RUNNING` implies that at least one worker of the current deployment instance's version has started running.
     ATLEAST_ONE_RUNNING = "at_least_one_running"
 
     # `ALL_RUNNING` implies that all workers of the current deployment instance's version have started running (i.e. all workers aligning to the minimum number of replicas).
     # It doesn't imply that all the workers relating to other deployments have been torn down.
     ALL_RUNNING = "all_running"
 
-    # `FULLY_FINISHED` implies Atleast min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
+    # `FULLY_FINISHED` implies at least min_replicas number of workers are running for the deployment and there are no pending or crashlooping workers from previous versions lying around.
     FULLY_FINISHED = "fully_finished"
 
     # `ASYNC` implies that the deployment will be assumed ready after the URL is minted and the worker statuses are not checked.
@@ -447,9 +447,9 @@ def _capsule_worker_semantic_status(
     crashlooping_workers = _make_version_dict(workers, "CrashLoopBackOff")
 
     # current_status (formulated basis):
-    # - atleast one pods are pending for `_end_state_capsule_version`
-    # - atleast one pod is in Running state for `_end_state_capsule_version` (maybe terminal) [Might require heath-check thing here]
-    # - alteast one pod is crashlooping for `_end_state_capsule_version` (maybe terminal)
+    # - at least one pods are pending for `_end_state_capsule_version`
+    # - at least one pod is in Running state for `_end_state_capsule_version` (maybe terminal) [Might require health-check thing here]
+    # - at least one pod is crashlooping for `_end_state_capsule_version` (maybe terminal)
     # - all pods are running for `_end_state_capsule_version` that match the minimum number of replicas
     # - all pods are running for `_end_state_capsule_version` that match the maximum number of replicas and no other pods of older versions are running
     # - no pods relating to `_end_state_capsule_version` are pending/running/crashlooping

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/app_cli.py

@@ -239,7 +239,7 @@ def _bake_image(app_config: AppConfig, cache_dir: str, logger):
         baking_status.resolved_image,
     )
     app_config.set_state("python_path", baking_status.python_path)
-    logger("🐳 Using The Docker Image : %s" % app_config.get_state("image"))
+    logger("🐳 Using the docker image : %s" % app_config.get_state("image"))
 
 
 def print_table(data, headers):
@@ -374,7 +374,7 @@ def _package_necessary_things(app_config: AppConfig, logger):
     #      or is it relative to where the caller command is sitting. Ideally it should work
     #      like Kustomizations where its relative to where the yaml file sits for simplicity
     #      of understanding relationships between config files. Ideally users can pass the src_path
-    #      from the command line and that will aliviate any need to package any other directories for
+    #      from the command line and that will alleviate any need to package any other directories for
     #
 
     package_dir = app_config.get_state("packaging_directory")
@@ -395,7 +395,7 @@ def _package_necessary_things(app_config: AppConfig, logger):
     )
     app_config.set_state("code_package_url", package_url)
     app_config.set_state("code_package_key", package_key)
-    logger("💾 Code Package Saved to : %s" % app_config.get_state("code_package_url"))
+    logger("💾 Code package saved to : %s" % app_config.get_state("code_package_url"))
 
 
 @app.command(help="Deploy an app to the Outerbounds Platform.")
@@ -465,7 +465,7 @@ def deploy(
 
         app_config.set_state("packaging_directory", packaging_directory)
         logger(
-            "📦 Packaging Directory : %s" % app_config.get_state("packaging_directory"),
+            "📦 Packaging directory : %s" % app_config.get_state("packaging_directory"),
         )
 
         if app_config.get("no_deps", False):
@@ -611,7 +611,7 @@ def deploy(
                     )
                 raise AppConfigError(message)
             capsule_logger(
-                f"🚀 {'' if not force_upgrade else 'Force'} Upgrading {capsule.capsule_type.lower()} `{capsule.name}`....",
+                f"🚀 {'Upgrading' if not force_upgrade else 'Force upgrading'} {capsule.capsule_type.lower()} `{capsule.name}`....",
                 color=ColorTheme.INFO_COLOR,
                 system_msg=True,
             )
@@ -632,7 +632,7 @@ def deploy(
             capsule_spinner.stop()
 
         logger(
-            f"💊 {capsule.capsule_type} {app_config.config['name']} ({capsule.identifier}) deployed! {capsule.capsule_type} exposed on the URL: {capsule.url}",
+            f"💊 {capsule.capsule_type} {app_config.config['name']} ({capsule.identifier}) deployed! {capsule.capsule_type} available on the URL: {capsule.url}",
             color=ColorTheme.INFO_COLOR,
             system_msg=True,
         )
@@ -761,7 +761,7 @@ def list(ctx, project, branch, name, tags, format, auth_type):
 def delete(ctx, name, cap_id, project, branch, tags, auto_approve):
 
     """Delete an app/apps from the Outerbounds Platform."""
-    # Atleast one of the args need to be provided
+    # At least one of the args need to be provided
     if not any(
         [
             name is not None,
@@ -772,7 +772,7 @@ def delete(ctx, name, cap_id, project, branch, tags, auto_approve):
         ]
     ):
         raise AppConfigError(
-            "Atleast one of the options need to be provided. You can use --name, --id, --project, --branch, --tag"
+            "At least one of the options need to be provided. You can use --name, --id, --project, --branch, --tag"
         )
 
     capsule_api = CapsuleApi(ctx.obj.api_url, ctx.obj.perimeter)

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/capsule.py

@@ -44,24 +44,24 @@ class CapsuleStateMachine:
     - Happy Path:
         - First time Create :
             - wait for status.updateInProgress to be set to False
-                - (interleved) Poll the worker endpoints to check their status
+                - (interleaved) Poll the worker endpoints to check their status
                     - showcase how many workers are coming up if things are on the cli side.
                 - If the user has set some flag like `--dont-wait-to-fully-finish` then we check the `status.currentlyServedVersion` to see if even one replica is ready to
                 serve traffic.
             - once the status.updateInProgress is set to False, it means that the replicas are ready
         - Upgrade:
             - wait for status.updateInProgress to be set to False
-                - (interleved) Poll the worker endpoints to check their status and signal the user the number replicas coming up
+                - (interleaved) Poll the worker endpoints to check their status and signal the user the number replicas coming up
                 - If the user has set some flag like `--dont-wait-to-fully-finish` then we check the `status.currentlyServedVersion` to see if even one replica is ready to
                 serve traffic.
     - Unhappy Path:
         - First time Create :
             - wait for status.updateInProgress to be set to False,
-                - (interleved) Poll the workers to check their status.
+                - (interleaved) Poll the workers to check their status.
                     - If the worker pertaining the current deployment instance version is crashlooping then crash the deployment process with the error messages and logs.
         - Upgrade:
             - wait for status.updateInProgress to be set to False,
-                - (interleved) Poll the workers to check their status.
+                - (interleaved) Poll the workers to check their status.
                     - If the worker pertaining the current deployment instance version is crashlooping then crash the deployment process with the error messages and logs.
 
     """
@@ -210,9 +210,9 @@ class CapsuleInput:
     def construct_exec_command(cls, commands: List[str]):
         commands = ["set -eEuo pipefail"] + commands
         command_string = "\n".join(commands)
-        # First constuct a base64 encoded string of the quoted command
+        # First construct a base64 encoded string of the quoted command
         # One of the reasons we don't directly pass the command string to the backend with a `\n` join
-        # is because the backend controller doesnt play nice when the command can be a multi-line string.
+        # is because the backend controller doesn't play nice when the command can be a multi-line string.
         # So we encode it to a base64 string and then decode it back to a command string at runtime to provide to
         # `bash -c`. The ideal thing to have done is to run "bash -c {shlex.quote(command_string)}" and call it a day
         # but the backend controller yields the following error:
@@ -682,7 +682,7 @@ class CapsuleDeployer:
         """
         - `capsule_response.version` contains the version of the object present in the database
         - `current_deployment_instance_version` contains the version of the object that was deployed by this instance of the deployer.
-        In the situtation that the versions of the objects become a mismatch then it means that current deployment process is not giving the user the
+        In the situation that the versions of the objects become a mismatch then it means that current deployment process is not giving the user the
         output that they desire.
         """
         if capsule_response.get("version", None) != current_deployment_instance_version:
@@ -783,24 +783,22 @@ class CapsuleDeployer:
             )
             if capsule_ready or failure_condition_satisfied:
                 logger(
-                    "💊 %s deployment status: %s | worker states: [success :%s | failure :%s ] "
+                    "💊 %s deployment status: %s "
                     % (
                         self.capsule_type.title(),
                         "in progress"
                         if state_machine.update_in_progress
                         else "completed",
-                        capsule_ready,
-                        failure_condition_satisfied,
                     )
                 )
                 _further_readiness_check_failed = False
                 if further_check_worker_readiness:
                     # HACK : monitor the workers for N seconds to make sure they are healthy
-                    # this is a hack. Ideally we should implment a healtcheck as a first class citizen
+                    # this is a hack. Ideally we should implement a healthcheck as a first class citizen
                     # but it will take some time to do that so in the meanwhile a timeout set on the cli
                     # side will be really helpful.
                     logger(
-                        "💊 running last minute readiness check for %s..."
+                        "💊 Running last minute readiness check for %s..."
                         % self.identifier
                     )
                     _further_readiness_check_failed = self._monitor_worker_readiness(

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/dependencies.py

@@ -95,7 +95,7 @@ def bake_deployment_image(
     pinned_conda_libs = get_pinned_conda_libs(python_version, DEFAULT_DATASTORE)
     pypi_packages.update(pinned_conda_libs)
     _reference = app_config.get("name", "default")
-    # `image` cannot be None. If by change it is none, FB will fart.
+    # `image` cannot be None. If by chance it is none, FB will fart.
     fb_response = bake_image(
         cache_file_path=cache_file_path,
         pypi_packages=pypi_packages,

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/secrets.py

@@ -116,7 +116,7 @@ class SecretRetriever:
 
         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."
+                "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:

{ob_metaflow_extensions-1.3.0 → ob_metaflow_extensions-1.3.1}/metaflow_extensions/outerbounds/plugins/apps/core/utils.py

@@ -185,8 +185,8 @@ def safe_requests_wrapper(
         - 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.
+        - 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.

ob_metaflow_extensions-1.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/__init__.py

@@ -0,0 +1,7 @@
+from .s3_proxy_decorator import (
+    S3ProxyDecorator,
+    NebiusS3ProxyDecorator,
+    CoreWeaveS3ProxyDecorator,
+)
+
+__all__ = ["S3ProxyDecorator", "NebiusS3ProxyDecorator", "CoreWeaveS3ProxyDecorator"]

ob_metaflow_extensions-1.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/constants.py

@@ -0,0 +1,8 @@
+S3_PROXY_BINARY_URLS = {
+    "aarch64": "https://fast-s3-proxy.outerbounds.sh/linux-arm64/s3-proxy-0.1.1.gz",
+    "x86_64": "https://fast-s3-proxy.outerbounds.sh/linux-amd64/s3-proxy-0.1.1.gz",
+}
+
+DEFAULT_PROXY_PORT = 8081
+DEFAULT_PROXY_HOST = "localhost"
+S3_PROXY_WRITE_MODES = ["origin-and-cache", "origin", "cache"]

ob_metaflow_extensions-1.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/exceptions.py

@@ -0,0 +1,13 @@
+from metaflow.exception import MetaflowException
+
+
+class S3ProxyException(MetaflowException):
+    headline = "S3 Proxy Error"
+
+
+class S3ProxyConfigException(S3ProxyException):
+    headline = "S3 Proxy Configuration Error"
+
+
+class S3ProxyApiException(S3ProxyException):
+    headline = "S3 Proxy API Error"

ob_metaflow_extensions-1.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_api.py

@@ -0,0 +1,93 @@
+import json
+import time
+from typing import Dict, Optional
+
+from .exceptions import S3ProxyConfigException, S3ProxyApiException
+
+
+class S3ProxyConfigResponse:
+    def __init__(self, data: Dict):
+        self.bucket_name = data.get("bucket_name")
+        self.endpoint_url = data.get("endpoint_url")
+        self.access_key_id = data.get("access_key_id")
+        self.secret_access_key = data.get("secret_access_key")
+        self.region = data.get("region")
+
+
+class S3ProxyApiClient:
+    def __init__(self):
+        self.perimeter, self.integrations_url = self._get_api_configs()
+
+    def _get_api_configs(self):
+        from metaflow_extensions.outerbounds.remote_config import init_config
+        from os import environ
+
+        conf = init_config()
+        perimeter = conf.get("OBP_PERIMETER") or environ.get("OBP_PERIMETER", "")
+        integrations_url = conf.get("OBP_INTEGRATIONS_URL") or environ.get(
+            "OBP_INTEGRATIONS_URL", ""
+        )
+
+        if not perimeter:
+            raise S3ProxyConfigException(
+                "No perimeter set. Please run `outerbounds configure` command."
+            )
+
+        if not integrations_url:
+            raise S3ProxyConfigException(
+                "No integrations URL set. Please contact your Outerbounds support team."
+            )
+
+        return perimeter, integrations_url
+
+    def fetch_s3_proxy_config(
+        self, integration_name: Optional[str] = None
+    ) -> S3ProxyConfigResponse:
+        url = f"{self.integrations_url}/s3proxy"
+
+        payload = {"perimeter_name": self.perimeter}
+        if integration_name:
+            payload["integration_name"] = integration_name
+
+        headers = {"Content-Type": "application/json"}
+
+        try:
+            from metaflow.metaflow_config import SERVICE_HEADERS
+
+            headers.update(SERVICE_HEADERS or {})
+        except ImportError:
+            pass
+
+        response = self._make_request(url, headers, payload)
+        return S3ProxyConfigResponse(response)
+
+    def _make_request(self, url: str, headers: Dict, payload: Dict) -> Dict:
+        from metaflow_extensions.outerbounds.plugins.secrets.secrets import (
+            _api_server_get,
+        )
+
+        retryable_status_codes = [409]
+        json_payload = json.dumps(payload)
+
+        for attempt in range(3):
+            response = _api_server_get(
+                url, data=json_payload, headers=headers, conn_error_retries=5
+            )
+
+            if response.status_code not in retryable_status_codes:
+                break
+
+            if attempt < 2:
+                time.sleep(0.5 * (attempt + 1))
+
+        if response.status_code != 200:
+            error_msg = f"API request failed with status {response.status_code}"
+            try:
+                error_data = response.json()
+                if "message" in error_data:
+                    error_msg = error_data["message"]
+            except:
+                pass
+            raise S3ProxyApiException(error_msg)
+
+        return response.json()

ob_metaflow_extensions-1.3.1/metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_decorator.py

@@ -0,0 +1,106 @@
+import functools
+from typing import Optional
+
+from metaflow import current
+from metaflow.decorators import StepDecorator
+
+from .s3_proxy_manager import S3ProxyManager
+from .exceptions import S3ProxyException
+from .constants import S3_PROXY_WRITE_MODES
+
+
+class S3ProxyDecorator(StepDecorator):
+    """
+    S3 Proxy decorator for routing S3 requests through a local proxy service.
+
+    Parameters
+    ----------
+    integration_name : str, optional
+        Name of the S3 proxy integration. If not specified, will use the only
+        available S3 proxy integration in the namespace (fails if multiple exist).
+    write_mode : str, optional
+        The desired behavior during write operations to target (origin) S3 bucket.
+        allowed options are:
+            "origin-and-cache" -> write to both the target S3 bucket and local object
+                storage
+            "origin" -> only write to the target S3 bucket
+            "cache" -> only write to the object storage service used for caching
+    debug : bool, optional
+        Enable debug logging for proxy operations.
+    """
+
+    name = "s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }
+
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        write_mode = self.attributes["write_mode"]
+        if write_mode and write_mode not in S3_PROXY_WRITE_MODES:
+            raise S3ProxyException(
+                f"unexpected write_mode specified: {write_mode}. Allowed values are: {','.join(S3_PROXY_WRITE_MODES)}."
+            )
+
+        self.manager = S3ProxyManager(
+            integration_name=self.attributes["integration_name"],
+            write_mode=self.attributes["write_mode"],
+            debug=self.attributes["debug"],
+        )
+
+        current._update_env({"s3_proxy": self.manager})
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        """Setup S3 proxy before step execution"""
+        self.manager.setup_proxy()
+
+    def task_finished(
+        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
+    ):
+        """Cleanup S3 proxy after step execution"""
+        if self.manager:
+            self.manager.cleanup()
+
+
+class NebiusS3ProxyDecorator(S3ProxyDecorator):
+    """
+    Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
+    It exists to make it easier for users to know that this decorator should only be used with
+    a Neo Cloud like Nebius.
+    """
+
+    name = "nebius_s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }
+
+
+class CoreWeaveS3ProxyDecorator(S3ProxyDecorator):
+    """
+    CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
+    It exists to make it easier for users to know that this decorator should only be used with
+    a Neo Cloud like CoreWeave.
+    """
+
+    name = "coreweave_s3_proxy"
+    defaults = {
+        "integration_name": None,
+        "write_mode": None,
+        "debug": False,
+    }