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

metaflow_extensions/outerbounds/plugins/apps/deploy_decorator.py

@@ -0,0 +1,201 @@
+from metaflow.exception import MetaflowException
+from metaflow.decorators import StepDecorator
+from metaflow import current
+from .app_utils import start_app
+from .supervisord_utils import SupervisorClient, SupervisorClientException
+import os
+import random
+import string
+import tempfile
+import sys
+from .consts import (
+    BASE_DIR_FOR_APP_ASSETS,
+    DEFAULT_WAIT_TIME_SECONDS_FOR_PROCESS_TO_START,
+)
+
+"""
+There are 3 variants of starting apps that we support through this function. Which variant is applied is
+a result of the user setting (or not setting) self.entrypoint and self.app_dir.
+
+The chosen variant determines whether or not the user will have an easy (auto-magical) way to write their
+Metaflow artifacts the the directory where their app will be run from. This will simplify managing models
+from inside the code.
+
+Case 1:
+    Desired Behavior:
+        The user doesn't care about auto-magical artifact management, they just have a server.py somewhere and
+        they want to run it.
+    How:
+        The user sets self.entrypoint to the name of their file (optionally with any args).
+        The value of self.app_dir is irrelevant.
+        Example: self.entrypoint = "/home/ob-workspace/my_random_directory/my_subfolder/server.py --my_arg 764"
+    Expected Behavior:
+        The users app will be started using a conda environment built by Metaflow, that's it.
+
+Case 2:
+    Desired Behavior:
+        The user has defined a clean package/module structure with a __main__.py in a folder somewhere. They would like
+        to run this module as an app. They would like to access models in the same top level directory as their app.
+    How:
+        The user sets self.entrypoint to None.
+        The user sets self.app_dir to the top level directory of their app.
+        Example: self.app_dir = "/home/ob-workspace/my_random_directory/my_subfolder"
+        (my_subfolder HAS to contain a __main__.py file.)
+    Expected Behavior:
+        The users app will be started using a conda environment built by Metaflow. The user gets access to self.deploy_dir where they
+        write their artifacts. After the user writes their artifacts to self.deploy_dir, the artifacts and the app is copied over
+        to an internal directory where we will deploy the app from.
+        The internal directory is: /home/ob-workspace/.appdaemon/apps/<app_port>.
+        <app_port> will contain user's __main__.py and other files (copied recursively), as well as their artifacts that they wrote to
+        self.deploy_dir. self.deploy_dir/my_model becomes <app_port>/my_model. my_subfolder/__main__.py becomes <app_port>/__main__.py.
+
+Case 3:
+    Desired Behavior:
+        The user follows the Outerbounds convention, which is the same as Case 2, except the app is actually in the same top level folder
+        as the Deployer flow. They would like to access models as usual.
+    How:
+        The user sets self.entrypoint to None (or doesn't set it at all).
+        The user sets self.app_dir to None (or doesn't set it at all).
+    Expected Behavior:
+        The users app will be started using a conda environment built by Metaflow. Everything else is exactly the same as Case 2.
+"""
+
+
+class WorkstationAppDeployDecorator(StepDecorator):
+    """
+    Specifies that this step is used to deploy an instance of the app.
+    Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
+
+    Parameters
+    ----------
+    app_port : int
+        Number of GPUs to use.
+    app_name : str
+        Name of the app to deploy.
+    """
+
+    name = "app_deploy"
+    defaults = {"app_port": 8080, "app_name": "app"}
+
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        if any([deco.name == "kubernetes" for deco in decos]):
+            raise MetaflowException(
+                "@app_deploy decorator is only supported locally and does not work with remote execution environments like @kubernetes, @nvidia."
+            )
+
+        # We always need to have some environment defined through the flow to deploy and app.
+        # Which means either step decorators like @pypi / @conda must be defined.
+        # or flow level decorators like @conda_base / @pypi_base.
+        if not any([deco.name == "pypi" or deco.name == "conda" for deco in decos]):
+            flow_decorators = flow._flow_decorators.keys()
+            if (
+                "conda_base" not in flow_decorators
+                and "pypi_base" not in flow_decorators
+            ):
+                raise MetaflowException(
+                    "@app_deploy requires either step decorators like @pypi / @conda or flow level decorators like @conda_base / @pypi_base to be defined."
+                )
+
+        app_port = self.attributes["app_port"]
+        app_name = self.attributes["app_name"]
+
+        # Currently this decorator is expected to only execute on workstation.
+        if app_port is None or app_port < 6000 or app_port > 6002:
+            raise MetaflowException(
+                "AppDeployDecorator requires app_port to be between 6000 and 6002."
+            )
+
+        if app_name is None:
+            raise MetaflowException("AppDeployDecorator requires app_name to be set.")
+
+    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,
+    ):
+        """
+        Runs before the step decorated with @app_deploy has started.
+        We create a directory where the user can write their artifacts and expose it via self.deploy_dir.
+        """
+        os.makedirs(BASE_DIR_FOR_APP_ASSETS, exist_ok=True)
+        # First we want to create a directory where the user's app directory and artifacts can be stored.
+        with tempfile.TemporaryDirectory(
+            prefix=BASE_DIR_FOR_APP_ASSETS, delete=False
+        ) as temp_dir:
+            launch_temp_dir = temp_dir
+
+        # Expose this to the user, so that they can use it write their artifacts.
+        setattr(flow, "deploy_dir", launch_temp_dir)
+
+        # Make sure to record deploy_dir so that the user cannot accidentally override it.
+        self._deploy_dir = launch_temp_dir
+
+    def task_post_step(
+        self, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        """
+        Runs after the step decorated with @app_deploy has finished.
+        Based on the cases above, things that we care about are self.entrypoint and self.app_dir.
+        """
+
+        deploy_dir = self._deploy_dir
+
+        # By default we assume that the user has a __main__.py file in their app directory.
+        # They can always override this behavior.
+        user_provided_entrypoint = getattr(flow, "entrypoint", None)
+
+        if user_provided_entrypoint is not None and not isinstance(
+            user_provided_entrypoint, str
+        ):
+            raise MetaflowException(
+                f"@app_deploy requires entrypoint to be set to a string. The current value of entrypoint {user_provided_entrypoint} is not valid."
+            )
+
+        flow_directory = os.path.dirname(os.path.abspath(sys.argv[0]))
+
+        app_location = getattr(
+            flow, "app_dir", os.path.join(flow_directory, self.attributes["app_name"])
+        )
+
+        if user_provided_entrypoint is None and not os.path.exists(app_location):
+            raise MetaflowException(f"App directory {app_location} does not exist.")
+
+        wait_time_for_app_start = getattr(
+            flow,
+            "wait_time_for_app_start",
+            DEFAULT_WAIT_TIME_SECONDS_FOR_PROCESS_TO_START,
+        )
+
+        try:
+            supervisor_client = SupervisorClient(
+                wait_time_seconds_for_app_start=wait_time_for_app_start
+            )
+
+            # First, let's deploy the app.
+            start_app(
+                port=self.attributes["app_port"], name=self.attributes["app_name"]
+            )
+
+            # Now, let's add the app to supervisor.
+            supervisor_client.start_process_with_supervisord(
+                self.attributes["app_name"],
+                self.attributes["app_port"],
+                user_provided_entrypoint,
+                deploy_dir,
+                app_location,
+            )
+        except SupervisorClientException as e:
+            raise MetaflowException(str(e))
+        except Exception as e:
+            raise MetaflowException(
+                f"Failed to start {self.attributes['app_name']}! Cause: {str(e)}"
+            ) from e

metaflow_extensions/outerbounds/plugins/apps/supervisord_utils.py

@@ -0,0 +1,243 @@
+import argparse
+import os
+import configparser
+import tempfile
+import sys
+import subprocess
+from pathlib import Path
+import shutil
+from enum import Enum
+import time
+from .consts import BASE_DIR_FOR_APP_ASSETS, APP_DAEMON_WORKSTAION_PATH
+
+
+class SupervisorClientException(Exception):
+    pass
+
+
+class SupervisorClient:
+    """
+    A client for starting and stopping apps using supervisor.
+    """
+
+    def __init__(self, wait_time_seconds_for_app_start: int):
+        self.supervisor_conf_loc = os.environ.get("SUPERVISOR_CONF_PATH")
+
+        self.wait_time_seconds_for_app_start = wait_time_seconds_for_app_start
+        if self.supervisor_conf_loc is None or not os.path.exists(
+            self.supervisor_conf_loc
+        ):
+            raise SupervisorClientException(
+                "This workstation does not support deploying apps! Please reach out to Outerbounds for support."
+            )
+
+        self.metaflow_envs_persistent_path = os.environ.get(
+            "SUPERVISOR_PYTHON_ENVS_PATH"
+        )
+        if self.metaflow_envs_persistent_path is None:
+            raise SupervisorClientException(
+                "This workstation does not support deploying apps! Please reach out to Outerbounds for support."
+            )
+
+        # Check if supervisorctl is installed
+        if not shutil.which("supervisorctl"):
+            raise SupervisorClientException(
+                "This workstation does not support deploying apps! Please reach out to Outerbounds for support."
+            )
+
+    def _stop_existing_app_at_port(self, app_port):
+        supervisor_config = configparser.ConfigParser()
+        supervisor_config.read(self.supervisor_conf_loc)
+
+        for program in supervisor_config.sections():
+            if "obp_app_port" in supervisor_config[program]:
+                if supervisor_config[program]["obp_app_port"].strip() == str(app_port):
+                    res = subprocess.run(
+                        ["supervisorctl", "stop", program],
+                        stdout=subprocess.DEVNULL,
+                        stderr=subprocess.DEVNULL,
+                    )
+
+                    del supervisor_config[program]
+
+        with tempfile.NamedTemporaryFile(
+            "w", dir=os.path.dirname(self.supervisor_conf_loc), delete=False
+        ) as f:
+            supervisor_config.write(f)
+            tmp_file = f.name
+
+        os.rename(tmp_file, self.supervisor_conf_loc)
+
+    def _create_supervisor_conf_entry(
+        self, command, launch_directory, app_port, app_name
+    ):
+        entry = {
+            "command": command,
+            "directory": launch_directory,
+            "autostart": "true",
+            "autorestart": "true",
+            "obp_app_port": app_port,  # Record the app port for internal reference. This is not used by supervisor.
+        }
+
+        supervisor_config = configparser.ConfigParser()
+        supervisor_config.read(self.supervisor_conf_loc)
+
+        supervisor_config[f"program:{app_name}"] = entry
+
+        with tempfile.NamedTemporaryFile(
+            "w", dir=os.path.dirname(self.supervisor_conf_loc), delete=False
+        ) as f:
+            supervisor_config.write(f)
+            tmp_file = f.name
+
+        os.rename(tmp_file, self.supervisor_conf_loc)
+
+    def start_process_with_supervisord(
+        self,
+        app_name,
+        app_port,
+        user_provided_entrypoint,
+        deploy_dir=None,
+        app_dir=None,
+    ):
+        """
+        Starts the app using supervisor.
+
+        Args:
+            app_name: The name of the app to start.
+            app_port: The port to start the app on.
+            user_provided_entrypoint: The entrypoint to start the app with.
+            deploy_dir: The directory to copy the app to and deploy from.
+            app_dir: The directory to copy the app from.
+        """
+
+        entrypoint = user_provided_entrypoint
+        deploy_dir_for_port = os.path.join(BASE_DIR_FOR_APP_ASSETS, str(app_port))
+        launch_directory = (
+            BASE_DIR_FOR_APP_ASSETS
+            if entrypoint is None
+            else APP_DAEMON_WORKSTAION_PATH
+        )
+
+        # Stop any existing apps that are running on the same port.
+        self._stop_existing_app_at_port(app_port)
+
+        # This means the user has opted for either case 2 or case 3.
+        # Cases 2 and 3 are handled the same way, the only thing that differs is self.app_dir.
+        if user_provided_entrypoint is None:
+            # Copy the app_dir to the deploy_dir.
+            # This is also where all (if any) artifacts are written.
+            recursive_copy(app_dir, deploy_dir)
+
+            # Copy the entire deploy_dir to the port specific directory.
+            # Clear out anything that was there before (maybe a different app's assets)
+            if os.path.exists(deploy_dir_for_port):
+                shutil.rmtree(deploy_dir_for_port)
+
+            os.makedirs(deploy_dir_for_port)
+            recursive_copy(deploy_dir, deploy_dir_for_port)
+
+            # Apply default value
+            # We launch the module from BASE_DIR_FOR_APP_ASSETS, so when we provide the -m flag we don't need (and can't use) the full path.
+            # We just need to provide the port number (which is also the name of the folder where all app assets are stored)
+            entrypoint = f"-m {str(app_port)}"
+
+        # deploy_dir is meant to be temporary. No need to keep it around after everything has been copied over.
+        shutil.rmtree(deploy_dir)
+
+        # Metaflow by default generates the environment in /root/... (which is not persisted on workstations).
+        # Since the environment is fully self contained, we can copy it to a persistent location.
+        persistent_path_for_executable = (
+            self._persist_metaflow_generated_python_environment()
+        )
+
+        # This is the command used by supervisord to launch the app.
+        command = f"{persistent_path_for_executable} {entrypoint}"
+
+        self._create_supervisor_conf_entry(
+            command, launch_directory, app_port, app_name
+        )
+
+        # Execute supervisorctl reload
+        # Capture the exit code
+        exit_code = subprocess.run(
+            ["supervisorctl", "reload"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        ).returncode
+        if exit_code != 0:
+            raise SupervisorClientException(
+                "Failed to start app! Contact Outerbounds for support."
+            )
+
+        print(
+            f"Waiting for {self.wait_time_seconds_for_app_start} seconds for {app_name} to start..."
+        )
+        time.sleep(self.wait_time_seconds_for_app_start)
+
+        self._raise_on_bad_status(app_name, command)
+
+    def _get_launched_prcoess_status(self, app_name, debug_command):
+        """
+        Checks the status of the launched process. If the status is not RUNNING or STARTING, it raises an exception.
+        Possible statuses: RUNNING, STARTING, STOPPED, BACKOFF, STOPPING, EXITED, FATAL, UNKNOWN
+
+        Args:
+            app_name: The name of the app to check the status of.
+            debug_command: The command to run to debug the app.
+        """
+        status_cmd_output = subprocess.run(
+            ["supervisorctl", "status", app_name],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        ).stdout.decode("utf-8")
+
+        status_cmd_output_parts = [
+            x.strip() for x in status_cmd_output.split(" ") if x.strip()
+        ]
+
+        status_str = status_cmd_output_parts[1]
+
+        if not status_str == "RUNNING" and not status_str == "STARTING":
+            raise SupervisorClientException(
+                f"Failed to start {app_name}! Try running {debug_command} manually to debug."
+            )
+
+    def _persist_metaflow_generated_python_environment(self):
+        """
+        Persists the metaflow generated python environment to a persistent location.
+        The step already runs in the environment generated by Metaflow.
+        """
+        current_executable = sys.executable
+        environment_path = Path(current_executable).parent.parent
+
+        persistent_path_for_this_environment = os.path.join(
+            self.metaflow_envs_persistent_path,
+            environment_path.parent.name,
+            environment_path.name,
+        )
+
+        final_executable_path = os.path.join(
+            persistent_path_for_this_environment,
+            Path(current_executable).parent.name,
+            Path(current_executable).name,
+        )
+
+        if os.path.exists(final_executable_path):
+            return final_executable_path
+
+        os.makedirs(persistent_path_for_this_environment, exist_ok=True)
+
+        recursive_copy(environment_path, persistent_path_for_this_environment)
+
+        return final_executable_path
+
+
+def recursive_copy(src, dst):
+    for item in os.listdir(src):
+        s = os.path.join(src, item)
+        d = os.path.join(dst, item)
+        if os.path.isdir(s):
+            shutil.copytree(s, d, dirs_exist_ok=True)
+        else:
+            shutil.copy2(s, d)

metaflow_extensions/outerbounds/plugins/auth_server.py

@@ -1,6 +1,10 @@
 from urllib.parse import urlparse
 
 import requests
+import time
+import random
+import os
+import sys
 
 
 def read_mf_config():
@@ -35,12 +39,28 @@ def get_token_url_and_headers(url_path):
 def get_token(url_path):
     from metaflow.exception import MetaflowException
 
-    url, headers = get_token_url_and_headers(url_path)
-    try:
-        r = requests.get(url, headers=headers)
-        r.raise_for_status()
-        token_info = r.json()
-        return token_info
+    retry_delay = 1
+    max_retries = os.environ.get("METAFLOW_EXT_HTTP_MAX_RETRIES", "10")
+    attempt = 0
+    while attempt < int(max_retries):
+        attempt = attempt + 1
+        try:
+            url, headers = get_token_url_and_headers(url_path)
+            r = requests.get(url, headers=headers)
+            r.raise_for_status()
+            token_info = r.json()
+            return token_info
+        except requests.exceptions.HTTPError as e:
+            raise MetaflowException(repr(e))
+        except (
+            requests.exceptions.ConnectionError,
+            requests.exceptions.ReadTimeout,
+        ) as e:
+            # ConnectionErrors are generally temporary errors like DNS resolution failures,
+            # timeouts etc.
+            time.sleep(retry_delay)
+            retry_delay *= 2  # Double the delay for the next attempt
+            retry_delay += random.uniform(0, 1)  # Add jitter
+            retry_delay = min(retry_delay, 10)
 
-    except requests.exceptions.HTTPError as e:
-        raise MetaflowException(repr(e))
+    raise MetaflowException("error connecting to Outerbounds")

metaflow_extensions/outerbounds/plugins/aws/__init__.py

@@ -0,0 +1,4 @@
+from .assume_role_decorator import assume_role
+from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
+
+__mf_promote_submodules__ = ["plugins.aws"]

metaflow_extensions/outerbounds/plugins/aws/assume_role.py

@@ -0,0 +1,3 @@
+# Environment variable name used by the assume_role decorator to pass
+# the role ARN to AWS client functions
+OBP_ASSUME_ROLE_ARN_ENV_VAR = "OBP_ASSUME_ROLE_ARN"

metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py

@@ -0,0 +1,118 @@
+from metaflow.user_decorators.mutable_flow import MutableFlow
+from metaflow.user_decorators.mutable_step import MutableStep
+from metaflow.user_decorators.user_flow_decorator import FlowMutator
+from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
+
+
+class assume_role(FlowMutator):
+    """
+    Flow-level decorator for assuming AWS IAM roles.
+
+    When applied to a flow, all steps in the flow will automatically use the specified IAM role-arn
+    as their source principal.
+
+    Usage:
+    ------
+    @assume_role(role_arn="arn:aws:iam::123456789012:role/my-iam-role")
+    class MyFlow(FlowSpec):
+        @step
+        def start(self):
+            import boto3
+            client = boto3.client("dynamodb")  # Automatically uses the role in the flow decorator
+            self.next(self.end)
+
+        @step
+        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(
+                "`role_arn` keyword argument is required for the assume_role decorator"
+            )
+
+        if not self.role_arn.startswith("arn:aws:iam::"):
+            raise ValueError(
+                "`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.
+        It sets up environment variables that will be used by the AWS client
+        to automatically assume the specified role.
+        """
+        # 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}
+            for d in step.decorator_specs:
+                name, _, _, deco_kwargs = d
+                if name == "environment":
+                    _env_kwargs.update(deco_kwargs["vars"])
+                    _step_has_env_set = True
+
+            if _step_has_env_set:
+                # remove the environment decorator
+                step.remove_decorator("environment")
+
+            # add the environment decorator
+            step.add_decorator(
+                environment,
+                deco_kwargs=dict(vars=_env_kwargs),
+            )
+
+        # Set the role ARN as an environment variable that will be picked up
+        # by the get_aws_client function
+        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 (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)