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

metaflow_extensions/outerbounds/__init__.py

@@ -1,12 +1,6 @@
 import metaflow.metaflow_config_funcs
 
-from metaflow_extensions.outerbounds.remote_config import init_config
-
-from metaflow_extensions.outerbounds.plugins.perimeters import (
-    set_current_perimeter_config_url_in_environment,
-)
-
-set_current_perimeter_config_url_in_environment()
+from metaflow_extensions.outerbounds.remote_config import init_config, reload_config
 
 # we want to overide OSS Metaflow's initialization behavior with our own to support remote configs
 # we're reassigning the METAFLOW_CONFIG variable because all downstream settings rely on it and

metaflow_extensions/outerbounds/config/__init__.py

@@ -1,5 +1,40 @@
+from metaflow.metaflow_config import from_conf
+
 DEFAULT_AWS_CLIENT_PROVIDER = "obp"
 
 DEFAULT_AZURE_CLIENT_PROVIDER = "obp"
 
 DEFAULT_GCP_CLIENT_PROVIDER = "obp"
+
+
+###
+# On Demand Docker image build configuration
+###
+# Image builder service url
+FAST_BAKERY_URL = from_conf("FAST_BAKERY_URL", None)
+
+
+###
+# NVCF configuration
+###
+# Maximum number of consecutive heartbeats that can be missed.
+NVIDIA_HEARTBEAT_THRESHOLD = from_conf("NVIDIA_HEARTBEAT_THRESHOLD", "3")
+
+
+###
+# Snowpark configuration
+###
+# Snowflake account to use with the @snowpark decorator
+SNOWPARK_ACCOUNT = from_conf("SNOWPARK_ACCOUNT")
+# Snowflake user to use with the @snowpark decorator
+SNOWPARK_USER = from_conf("SNOWPARK_USER")
+# Snowflake password to use with the @snowpark decorator
+SNOWPARK_PASSWORD = from_conf("SNOWPARK_PASSWORD")
+# Snowflake role to use with the @snowpark decorator
+SNOWPARK_ROLE = from_conf("SNOWPARK_ROLE")
+# Snowflake database to use with the @snowpark decorator
+SNOWPARK_DATABASE = from_conf("SNOWPARK_DATABASE")
+# Snowflake warehouse to use with the @snowpark decorator
+SNOWPARK_WAREHOUSE = from_conf("SNOWPARK_WAREHOUSE")
+# Snowflake schema to use with the @snowpark decorator
+SNOWPARK_SCHEMA = from_conf("SNOWPARK_SCHEMA")

metaflow_extensions/outerbounds/plugins/__init__.py

@@ -32,6 +32,134 @@ def hide_access_keys(*args, **kwds):
             os.environ["AWS_SESSION_TOKEN"] = AWS_SESSION_TOKEN
 
 
+# This is a special placeholder value that can be passed as role_arn to
+# get_boto3_session() which makes it use the CSPR role, if its set.
+USE_CSPR_ROLE_ARN_IF_SET = "__cspr__"
+
+
+def get_boto3_session(role_arn=None, session_vars=None):
+    import boto3
+    import botocore
+    from metaflow_extensions.outerbounds.plugins.auth_server import get_token
+    from metaflow_extensions.outerbounds.plugins.aws.assume_role import (
+        OBP_ASSUME_ROLE_ARN_ENV_VAR,
+    )
+
+    from hashlib import sha256
+    from metaflow.util import get_username
+
+    user = get_username()
+
+    token_info = get_token("/generate/aws")
+
+    # Write token to a file. The file name is derived from the user name
+    # so it works with multiple users on the same machine.
+    #
+    # We hash the user name so we don't have to deal with special characters
+    # in the file name and the file name is not exposed to the user
+    # anyways, so it doesn't matter that its a little ugly.
+    token_file = "/tmp/obp_token." + sha256(user.encode("utf-8")).hexdigest()[:16]
+
+    # Write to a temp file then rename to avoid a situation when someone
+    # tries to read the file after it was open for writing (and truncated)
+    # but before the token was written to it.
+    with tempfile.NamedTemporaryFile("w", delete=False) as f:
+        f.write(token_info["token"])
+        tmp_token_file = f.name
+    os.rename(tmp_token_file, token_file)
+
+    cspr_role = None
+    if token_info.get("cspr_role_arn"):
+        cspr_role = token_info["cspr_role_arn"]
+
+    # Check if the assume_role decorator has set a CSPR ARN via environment variable
+    # This takes precedence over CSPR role that comes from the token_info response
+    decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+    if decorator_role_arn:
+        cspr_role = decorator_role_arn
+
+    if cspr_role:
+        # If CSPR role is set, we set it as the default role to assume
+        # for the AWS SDK. We do this by writing an AWS config file
+        # with two profiles. One to get credentials for the task role
+        # in exchange for the OIDC token, and second to assume the
+        # CSPR role using the task role credentials.
+        import configparser
+        from io import StringIO
+
+        aws_config = configparser.ConfigParser()
+
+        # Task role profile
+        aws_config["profile task"] = {
+            "role_arn": token_info["role_arn"],
+            "web_identity_token_file": token_file,
+        }
+
+        # CSPR role profile (default)
+        aws_config["profile cspr"] = {
+            "role_arn": cspr_role,
+            "source_profile": "task",
+        }
+
+        aws_config_string = StringIO()
+        aws_config.write(aws_config_string)
+        aws_config_file = (
+            "/tmp/aws_config." + sha256(user.encode("utf-8")).hexdigest()[:16]
+        )
+        with tempfile.NamedTemporaryFile(
+            "w", delete=False, dir=os.path.dirname(aws_config_file)
+        ) as f:
+            f.write(aws_config_string.getvalue())
+            tmp_aws_config_file = f.name
+        os.rename(tmp_aws_config_file, aws_config_file)
+        os.environ["AWS_CONFIG_FILE"] = aws_config_file
+        os.environ["AWS_PROFILE"] = "cspr"
+    else:
+        os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
+        os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
+
+    # Enable regional STS endpoints. This is the new recommended way
+    # by AWS [1] and is the more performant way.
+    # [1] https://docs.aws.amazon.com/sdkref/latest/guide/feature-sts-regionalized-endpoints.html
+    os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional"
+    if token_info.get("region"):
+        os.environ["AWS_DEFAULT_REGION"] = token_info["region"]
+
+    if cspr_role:
+        #  The generated AWS config will be used here since we set the
+        #  AWS_CONFIG_FILE environment variable above.
+        if role_arn == USE_CSPR_ROLE_ARN_IF_SET:
+            # Otherwise start from the default profile, assuming CSPR role
+            session = boto3.session.Session(profile_name="cspr")
+        else:
+            session = boto3.session.Session(profile_name="task")
+    else:
+        # Not using AWS config, just AWS_WEB_IDENTITY_TOKEN_FILE + AWS_ROLE_ARN
+        session = boto3.session.Session()
+
+    if role_arn and role_arn != USE_CSPR_ROLE_ARN_IF_SET:
+        # If the user provided a role_arn, we assume that role
+        # using the task role credentials. CSPR role is not used.
+        fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
+            client_creator=session._session.create_client,
+            source_credentials=session._session.get_credentials(),
+            role_arn=role_arn,
+            extra_args={},
+        )
+        creds = botocore.credentials.DeferredRefreshableCredentials(
+            method="assume-role", refresh_using=fetcher.fetch_credentials
+        )
+        botocore_session = botocore.session.Session(session_vars=session_vars)
+        botocore_session._credentials = creds
+        return boto3.session.Session(botocore_session=botocore_session)
+    else:
+        # If the user didn't provide a role_arn, or if the role_arn
+        # is set to USE_CSPR_ROLE_ARN_IF_SET, we return the default
+        # session which would use the CSPR role if it is set on the
+        # server, and the task role otherwise.
+        return session
+
+
 class ObpAuthProvider(object):
     name = "obp"
 
@@ -42,67 +170,19 @@ class ObpAuthProvider(object):
         if client_params is None:
             client_params = {}
 
-        import boto3
-        import botocore
         from botocore.exceptions import ClientError
-        from metaflow_extensions.outerbounds.plugins.auth_server import get_token
-
-        from hashlib import sha256
-        from metaflow.util import get_username
-
-        user = get_username()
-
-        token_info = get_token("/generate/aws")
-
-        # Write token to a file. The file name is derived from the user name
-        # so it works with multiple users on the same machine.
-        #
-        # We hash the user name so we don't have to deal with special characters
-        # in the file name and the file name is not exposed to the user
-        # anyways, so it doesn't matter that its a little ugly.
-        token_file = "/tmp/obp_token." + sha256(user.encode("utf-8")).hexdigest()[:16]
-
-        # Write to a temp file then rename to avoid a situation when someone
-        # tries to read the file after it was open for writing (and truncated)
-        # but before the token was written to it.
-        with tempfile.NamedTemporaryFile("w", delete=False) as f:
-            f.write(token_info["token"])
-            tmp_token_file = f.name
-        os.rename(tmp_token_file, token_file)
-
-        os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
-        os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
-
-        # Enable regional STS endpoints. This is the new recommended way
-        # by AWS [1] and is the more performant way.
-        # [1] https://docs.aws.amazon.com/sdkref/latest/guide/feature-sts-regionalized-endpoints.html
-        os.environ["AWS_STS_REGIONAL_ENDPOINTS"] = "regional"
-        if token_info.get("region"):
-            os.environ["AWS_DEFAULT_REGION"] = token_info["region"]
+        from botocore.config import Config
 
         with hide_access_keys():
-            if role_arn:
-                session = boto3.session.Session()
-                fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
-                    client_creator=session._session.create_client,
-                    source_credentials=session._session.get_credentials(),
-                    role_arn=role_arn,
-                    extra_args={},
-                )
-                creds = botocore.credentials.DeferredRefreshableCredentials(
-                    method="assume-role", refresh_using=fetcher.fetch_credentials
-                )
-                botocore_session = botocore.session.Session(session_vars=session_vars)
-                botocore_session._credentials = creds
-                session = boto3.session.Session(botocore_session=botocore_session)
-                if with_error:
-                    return session.client(module, **client_params), ClientError
-                else:
-                    return session.client(module, **client_params)
+            session = get_boto3_session(role_arn, session_vars)
+            _client_params = client_params.copy()
+            if _client_params.get("config") and type(_client_params["config"]) == dict:
+                _client_params["config"] = Config(**_client_params["config"])
+
             if with_error:
-                return boto3.client(module, **client_params), ClientError
+                return session.client(module, **_client_params), ClientError
             else:
-                return boto3.client(module, **client_params)
+                return session.client(module, **_client_params)
 
 
 AWS_CLIENT_PROVIDERS_DESC = [("obp", ".ObpAuthProvider")]
@@ -178,7 +258,6 @@ class ObpGcpAuthProvider(object):
 
     @staticmethod
     def get_gs_storage_client(*args, **kwargs):
-
         import sys
         from metaflow_extensions.outerbounds.plugins.auth_server import get_token
 
@@ -240,3 +319,53 @@ class ObpGcpAuthProvider(object):
 
 
 GCP_CLIENT_PROVIDERS_DESC = [("obp", ".ObpGcpAuthProvider")]
+CLIS_DESC = [
+    ("nvidia", ".nvcf.nvcf_cli.cli"),
+    ("nvct", ".nvct.nvct_cli.cli"),
+    ("fast-bakery", ".fast_bakery.fast_bakery_cli.cli"),
+    ("snowpark", ".snowpark.snowpark_cli.cli"),
+]
+STEP_DECORATORS_DESC = [
+    ("nvidia", ".nvcf.nvcf_decorator.NvcfDecorator"),
+    ("nvct", ".nvct.nvct_decorator.NvctDecorator"),
+    (
+        "fast_bakery_internal",
+        ".fast_bakery.fast_bakery_decorator.InternalFastBakeryDecorator",
+    ),
+    ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
+    ("tensorboard", ".tensorboard.TensorboardDecorator"),
+    ("gpu_profile", ".profilers.gpu_profile_decorator.GPUProfileDecorator"),
+    ("test_append_card", ".profilers.simple_card_decorator.DynamicCardAppendDecorator"),
+    ("nim", ".nim.nim_decorator.NimDecorator"),
+    ("ollama", ".ollama.OllamaDecorator"),
+    ("vllm", ".vllm.VLLMDecorator"),
+    ("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"),
+    (
+        "app_deploy_internal",
+        ".apps.core.app_deploy_decorator.AppDeployInternalDecorator",
+    ),
+]
+
+FLOW_DECORATORS_DESC = [
+    ("app_deploy", ".apps.core.app_deploy_decorator.AppDeployFlowDecorator"),
+]
+
+TOGGLE_STEP_DECORATOR = [
+    "-batch",
+    "-step_functions_internal",
+    "-airflow_internal",
+]
+
+TOGGLE_CLI = ["-batch", "-step-functions", "-airflow"]
+
+ENVIRONMENTS_DESC = [
+    ("fast-bakery", ".fast_bakery.docker_environment.DockerEnvironment")
+]
+
+SECRETS_PROVIDERS_DESC = [
+    ("outerbounds", ".secrets.secrets.OuterboundsSecretsProvider"),
+]
+# Adding an override here so the library can be imported at the metaflow.plugins level
+__mf_promote_submodules__ = ["snowflake", "ollama", "torchtune", "optuna"]

metaflow_extensions/outerbounds/plugins/apps/app_utils.py

@@ -0,0 +1,187 @@
+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=""):
+    """
+    Starts an app on the workstation.
+    List workstations, looks for "NamedPorts", then makes an update call to the NamedPorts for the workstation.
+    """
+    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."
+        )
+
+    # Every workstation has this environment variable set.
+    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):
+    """
+    Ensures that the port number is available on the workstation and that an app of
+    the same name is not already opened on a different port.
+
+    Args:
+        existing_named_ports: A list of named ports on the workstation.
+        port: The port number to check.
+        name: The name of the app to check.
+    """
+    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:
+    """
+    Waits for the app to be ready by polling the workstation status.
+    """
+    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/consts.py

@@ -0,0 +1,3 @@
+DEFAULT_WAIT_TIME_SECONDS_FOR_PROCESS_TO_START = 10
+BASE_DIR_FOR_APP_ASSETS = "/home/ob-workspace/.appdaemon/apps/"
+APP_DAEMON_WORKSTAION_PATH = "/home/ob-workspace/.appdaemon"

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

@@ -0,0 +1,15 @@
+from . import config
+from . import dependencies
+from . import capsule
+from . import utils
+from . import app_config
+from . import code_package
+from .deployer import AppDeployer, bake_image, package_code, DeployedApp
+from .config import BakedImage, PackagedCode
+from .config.typed_configs import (
+    ReplicaConfigDict,
+    ResourceConfigDict,
+    AuthConfigDict,
+    DependencyConfigDict,
+    PackageConfigDict,
+)