ob-metaflow-extensions 1.1.122__py2.py3-none-any.whl → 1.1.123rc0__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -319,6 +319,7 @@ STEP_DECORATORS_DESC = [
     ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
     ("tensorboard", ".tensorboard.TensorboardDecorator"),
     ("gpu_profile", ".profilers.gpu_profile_decorator.GPUProfileDecorator"),
+    ("app_deploy", ".apps.deploy_decorator.WorkstationAppDeployDecorator"),
 ]
 FLOW_DECORATORS_DESC = [("nim", ".nim.NimDecorator")]
 TOGGLE_STEP_DECORATOR = [

metaflow_extensions/outerbounds/plugins/apps/app_utils.py

@@ -0,0 +1,170 @@
+from metaflow.exception import MetaflowException
+import os
+from metaflow.metaflow_config_funcs import init_config
+import requests
+import time
+import random
+
+# IMPORTANT: Currently contents of this file are mostly duplicated from the outerbounds package.
+# This is purely due to the time rush of having to deliver this feature. As a fast forward, we
+# will reorganize things in a way that the amount of duplication in minimum.
+
+
+APP_READY_POLL_TIMEOUT_SECONDS = 300
+# Even after our backend validates that the app routes are ready, it takes a few seconds for
+# the app to be accessible via the browser. Till we hunt down this delay, add an extra buffer.
+APP_READY_EXTRA_BUFFER_SECONDS = 30
+
+
+def start_app(port=-1, name=""):
+    if len(name) == 0 or len(name) >= 20:
+        raise MetaflowException("App name should not be more than 20 characters long.")
+    elif not name.isalnum() or not name.islower():
+        raise MetaflowException(
+            "App name can only contain lowercase alphanumeric characters."
+        )
+
+    if "WORKSTATION_ID" not in os.environ:
+        raise MetaflowException(
+            "All outerbounds app commands can only be run from a workstation."
+        )
+
+    workstation_id = os.environ["WORKSTATION_ID"]
+
+    try:
+        try:
+            conf = init_config()
+            metaflow_token = conf["METAFLOW_SERVICE_AUTH_KEY"]
+            api_url = conf["OBP_API_SERVER"]
+
+            workstations_response = requests.get(
+                f"https://{api_url}/v1/workstations",
+                headers={"x-api-key": metaflow_token},
+            )
+            workstations_response.raise_for_status()
+        except:
+            raise MetaflowException("Failed to list workstations!")
+
+        workstations_json = workstations_response.json()["workstations"]
+        for workstation in workstations_json:
+            if workstation["instance_id"] == os.environ["WORKSTATION_ID"]:
+                if "named_ports" in workstation["spec"]:
+                    try:
+                        ensure_app_start_request_is_valid(
+                            workstation["spec"]["named_ports"], port, name
+                        )
+                    except ValueError as e:
+                        raise MetaflowException(str(e))
+
+                    for named_port in workstation["spec"]["named_ports"]:
+                        if int(named_port["port"]) == port:
+                            if named_port["enabled"] and named_port["name"] == name:
+                                print(f"App {name} started on port {port}!")
+                                print(
+                                    f"Browser URL: https://{api_url.replace('api', 'ui')}/apps/{os.environ['WORKSTATION_ID']}/{name}/"
+                                )
+                                print(
+                                    f"API URL: https://{api_url}/apps/{os.environ['WORKSTATION_ID']}/{name}/"
+                                )
+                                return
+                            else:
+                                try:
+                                    response = requests.put(
+                                        f"https://{api_url}/v1/workstations/update/{workstation_id}/namedports",
+                                        headers={"x-api-key": metaflow_token},
+                                        json={
+                                            "port": port,
+                                            "name": name,
+                                            "enabled": True,
+                                        },
+                                    )
+
+                                    response.raise_for_status()
+                                    poll_success = wait_for_app_port_to_be_accessible(
+                                        api_url,
+                                        metaflow_token,
+                                        workstation_id,
+                                        name,
+                                        APP_READY_POLL_TIMEOUT_SECONDS,
+                                    )
+                                    if poll_success:
+                                        print(f"App {name} started on port {port}!")
+                                        print(
+                                            f"Browser URL: https://{api_url.replace('api', 'ui')}/apps/{os.environ['WORKSTATION_ID']}/{name}/"
+                                        )
+                                        print(
+                                            f"API URL: https://{api_url}/apps/{os.environ['WORKSTATION_ID']}/{name}/"
+                                        )
+                                    else:
+                                        raise MetaflowException(
+                                            f"The app could not be deployed in {APP_READY_POLL_TIMEOUT_SECONDS / 60} minutes. Please try again later."
+                                        )
+                                except Exception:
+                                    raise MetaflowException(
+                                        f"Failed to start app {name} on port {port}!"
+                                    )
+    except Exception as e:
+        raise MetaflowException(f"Failed to start app {name} on port {port}!")
+
+
+def ensure_app_start_request_is_valid(existing_named_ports, port: int, name: str):
+    existing_apps_by_port = {np["port"]: np for np in existing_named_ports}
+
+    if port not in existing_apps_by_port:
+        raise MetaflowException(f"Port {port} not found on workstation")
+
+    for existing_named_port in existing_named_ports:
+        if (
+            name == existing_named_port["name"]
+            and existing_named_port["port"] != port
+            and existing_named_port["enabled"]
+        ):
+            raise MetaflowException(
+                f"App with name '{name}' is already deployed on port {existing_named_port['port']}"
+            )
+
+
+def wait_for_app_port_to_be_accessible(
+    api_url, metaflow_token, workstation_id, app_name, poll_timeout_seconds
+) -> bool:
+    num_retries_per_request = 3
+    start_time = time.time()
+    retry_delay = 1.0
+    poll_interval = 10
+    wait_message = f"App {app_name} is currently being deployed..."
+    while time.time() - start_time < poll_timeout_seconds:
+        for _ in range(num_retries_per_request):
+            try:
+                workstations_response = requests.get(
+                    f"https://{api_url}/v1/workstations",
+                    headers={"x-api-key": metaflow_token},
+                )
+                workstations_response.raise_for_status()
+                if is_app_ready(workstations_response.json(), workstation_id, app_name):
+                    print(wait_message)
+                    time.sleep(APP_READY_EXTRA_BUFFER_SECONDS)
+                    return True
+                else:
+                    print(wait_message)
+                    time.sleep(poll_interval)
+            except (
+                requests.exceptions.ConnectionError,
+                requests.exceptions.ReadTimeout,
+            ):
+                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)
+    return False
+
+
+def is_app_ready(response_json: dict, workstation_id: str, app_name: str) -> bool:
+    """Checks if the app is ready in the given workstation's response."""
+    workstations = response_json.get("workstations", [])
+    for workstation in workstations:
+        if workstation.get("instance_id") == workstation_id:
+            hosted_apps = workstation.get("status", {}).get("hosted_apps", [])
+            for hosted_app in hosted_apps:
+                if hosted_app.get("name") == app_name:
+                    return bool(hosted_app.get("ready"))
+    return False

metaflow_extensions/outerbounds/plugins/apps/deploy_decorator.py

@@ -0,0 +1,90 @@
+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
+
+DEFAULT_WAIT_TIME_SECONDS_FOR_PROCESS_TO_START = 10
+
+
+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(
+                "Step *{step}* is marked for execution both on Kubernetes and "
+                "Nvidia. Please use one or the other.".format(step=step)
+            )
+
+        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_post_step(
+        self, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        entrypoint = getattr(flow, "entrypoint", None)
+        wait_time_for_app_start = getattr(
+            flow,
+            "wait_time_for_app_start",
+            DEFAULT_WAIT_TIME_SECONDS_FOR_PROCESS_TO_START,
+        )
+
+        if entrypoint is None or not isinstance(entrypoint, str):
+            raise MetaflowException(
+                f"@app_deploy requires entrypoint to be set to a string. The current value of entrypoint {entrypoint} is not valid."
+            )
+
+        launch_dir = getattr(flow, "launch_dir", None)
+        if launch_dir is None or launch_dir == "":
+            raise MetaflowException("@app_deploy requires launch_dir to be set.")
+        elif not isinstance(launch_dir, str) or not os.path.exists(launch_dir):
+            raise MetaflowException(
+                f"@app_deploy requires launch_dir to be set to a valid directory. The current value of launch_dir {launch_dir} is not valid."
+            )
+
+        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"],
+                entrypoint,
+                launch_dir,
+            )
+        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,167 @@
+import argparse
+import os
+import configparser
+import tempfile
+import sys
+import subprocess
+from pathlib import Path
+import shutil
+from enum import Enum
+import time
+
+
+class SupervisorClientException(Exception):
+    pass
+
+
+class SupervisorClient:
+    class SupervisodProcessCodes(Enum):
+        STOPPED = 0
+        STARTING = 10
+        RUNNING = 20
+        BACKOFF = 30
+        STOPPING = 40
+        EXITED = 100
+        FATAL = 200
+        UNKNOWN = 1000
+
+    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 start_process_with_supervisord(self, app_name, entrypoint, directory=None):
+        """Add a new program entry to supervisor configuration."""
+
+        persistent_path_for_executable = (
+            self.persist_metaflow_generated_python_environment()
+        )
+
+        command = f"{persistent_path_for_executable} {entrypoint}"
+
+        entry = {
+            "command": command,
+            "directory": directory or os.getcwd(),
+            "autostart": "true",
+            "autorestart": "true",
+        }
+
+        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)
+
+        # Execute supervisorctl reload
+        # Capture the exit code
+        exit_code = subprocess.run(
+            ["supervisorctl", "reload"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        ).returncode
+        if exit_code != 0:
+            print("Failed to reload supervisor configuration!", file=sys.stderr)
+            return
+
+        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)
+        status = self._get_launched_prcoess_status(app_name)
+
+        if status not in [
+            self.SupervisodProcessCodes.RUNNING,
+            self.SupervisodProcessCodes.STARTING,
+        ]:
+            raise SupervisorClientException(
+                f"Failed to start {app_name}! Try running {command} manually to debug."
+            )
+
+    def _get_launched_prcoess_status(self, app_name):
+        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 status_str == "RUNNING":
+            return self.SupervisodProcessCodes.RUNNING
+        elif status_str == "STOPPED":
+            return self.SupervisodProcessCodes.STOPPED
+        elif status_str == "STARTING":
+            return self.SupervisodProcessCodes.STARTING
+        elif status_str == "BACKOFF":
+            return self.SupervisodProcessCodes.BACKOFF
+        elif status_str == "STOPPING":
+            return self.SupervisodProcessCodes.STOPPING
+        elif status_str == "EXITED":
+            return self.SupervisodProcessCodes.EXITED
+        elif status_str == "FATAL":
+            return self.SupervisodProcessCodes.FATAL
+        else:
+            return self.SupervisodProcessCodes.UNKNOWN
+
+    # By default, an environment generated by metaflow will end up in a path like: /root/micromamba/envs/metaflow/linux-64/02699a4d2d50cfc/bin/python
+    # However, on a workstation these environments are not persisted, so we need to copy them over to /home/ob-workspace
+    def persist_metaflow_generated_python_environment(self):
+        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)
+
+        for item in os.listdir(environment_path):
+            src_path = os.path.join(environment_path, item)
+            dst_path = os.path.join(persistent_path_for_this_environment, item)
+            if os.path.isdir(src_path):
+                shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
+            else:
+                shutil.copy2(src_path, dst_path)
+
+        return final_executable_path

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py

@@ -207,7 +207,7 @@ class RunningJob(object):
 
     @property
     def is_waiting(self):
-        return self.status == "PENDING"
+        return self.status in ["PENDING", "UNKNOWN"]
 
     @property
     def is_running(self):

{ob_metaflow_extensions-1.1.122.dist-info → ob_metaflow_extensions-1.1.123rc0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.1.122
+Version: 1.1.123rc0
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial

{ob_metaflow_extensions-1.1.122.dist-info → ob_metaflow_extensions-1.1.123rc0.dist-info}/RECORD

@@ -1,9 +1,13 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=TRGvIUMjkfneWtYUFSWoubu_Kf2ekAL4WLbV3IxOj9k,499
 metaflow_extensions/outerbounds/remote_config.py,sha256=Zpfpjgz68_ZgxlXezjzlsDLo4840rkWuZgwDB_5H57U,4059
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=QDcoAAiLEWG1-_mZhpswwtgtCZJvyGEsZOhNau7S9L4,12768
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=0M_MbFmDcHRYs9mJ7SFDT1hmXkVEhzEEWRgUU7mEt2w,12844
 metaflow_extensions/outerbounds/plugins/auth_server.py,sha256=_Q9_2EL0Xy77bCRphkwT1aSu8gQXRDOH-Z-RxTUO8N4,2202
 metaflow_extensions/outerbounds/plugins/perimeters.py,sha256=QXh3SFP7GQbS-RAIxUOPbhPzQ7KDFVxZkTdKqFKgXjI,2697
+metaflow_extensions/outerbounds/plugins/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+metaflow_extensions/outerbounds/plugins/apps/app_utils.py,sha256=JrVKlbRx8-nSmI4cRrB7F8BQGDHleABIYZudK4P-XFE,7905
+metaflow_extensions/outerbounds/plugins/apps/deploy_decorator.py,sha256=goD50XtFOrTl8jsHazdYbu6owVzReOL_1YeukOgX75E,3439
+metaflow_extensions/outerbounds/plugins/apps/supervisord_utils.py,sha256=1ImWHiQ4GNIyognIg-IT-cpInS1-CKWCEy2h9upZufo,5975
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=i6F3FXwvEhkmUCTHDJ4VmSoL6vKyQhC_YRCtY6F4EkA,14209
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=cj63FrdioggipQFP8GwgxU3FYe6IyzjGSUGYxLQZ4nQ,5189
@@ -33,7 +37,7 @@ metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py,sha256=ezJ2Jr8J
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py,sha256=JEW0EUxj_mNZXo9OFkJFmWfg-P7_CEgvNbgsMTCBTAE,4273
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py,sha256=a7LqSKULVh8IrR1StrVPbemHOLojR0nEqh-mMX-M1i4,9904
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_exceptions.py,sha256=FTfYlJu-sn9DkPOs2R1V1ChWb1vZthOgeq0BZdT1ucY,296
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=d_5UhXqZ_12rCvatH1capPQZYGLx1FVqq_rtW65OXyk,6874
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=aQphxX6jqYgfa83w387pEWl0keuLm38V53I8P8UL2ck,6887
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_service_spec.py,sha256=AI_kcm1hZV3JRxJkookcH6twiGnAYjk9Dx-MeoYz60Y,8511
 metaflow_extensions/outerbounds/plugins/tensorboard/__init__.py,sha256=9lUM4Cqi5RjrHBRfG6AQMRz8-R96eZC8Ih0KD2lv22Y,1858
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
@@ -44,7 +48,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2
 metaflow_extensions/outerbounds/toplevel/plugins/gcp/__init__.py,sha256=BbZiaH3uILlEZ6ntBLKeNyqn3If8nIXZFq_Apd7Dhco,70
 metaflow_extensions/outerbounds/toplevel/plugins/kubernetes/__init__.py,sha256=5zG8gShSj8m7rgF4xgWBZFuY3GDP5n1T0ktjRpGJLHA,69
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
-ob_metaflow_extensions-1.1.122.dist-info/METADATA,sha256=cSdliiwbdi-0Cr7GI4nuQfRQyGfmdrXuU3BPR3UJQf0,520
-ob_metaflow_extensions-1.1.122.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.122.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.122.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.123rc0.dist-info/METADATA,sha256=jVJeDkUJvaWZ3B2Xkor72wjgNKwd1u3LThYzdotjQYs,523
+ob_metaflow_extensions-1.1.123rc0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.123rc0.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.123rc0.dist-info/RECORD,,