ob-metaflow-extensions 1.1.151__py2.py3-none-any.whl → 1.4.33__py2.py3-none-any.whl

metaflow_extensions/outerbounds/__init__.py

@@ -1,6 +1,6 @@
 import metaflow.metaflow_config_funcs
 
-from metaflow_extensions.outerbounds.remote_config import init_config
+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/plugins/__init__.py

@@ -41,6 +41,9 @@ 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
@@ -69,6 +72,12 @@ def get_boto3_session(role_arn=None, session_vars=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
@@ -326,9 +335,14 @@ STEP_DECORATORS_DESC = [
     ("snowpark", ".snowpark.snowpark_decorator.SnowparkDecorator"),
     ("tensorboard", ".tensorboard.TensorboardDecorator"),
     ("gpu_profile", ".profilers.gpu_profile_decorator.GPUProfileDecorator"),
-    ("nim", ".nim.NimDecorator"),
+    ("test_append_card", ".profilers.simple_card_decorator.DynamicCardAppendDecorator"),
+    ("nim", ".nim.nim_decorator.NimDecorator"),
     ("ollama", ".ollama.OllamaDecorator"),
-    ("app_deploy", ".apps.deploy_decorator.WorkstationAppDeployDecorator"),
+    ("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 = [
@@ -347,4 +361,4 @@ 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"]
+__mf_promote_submodules__ = ["snowflake", "ollama", "torchtune", "optuna"]

metaflow_extensions/outerbounds/plugins/apps/app_deploy_decorator.py

@@ -0,0 +1,146 @@
+from metaflow.exception import MetaflowException
+from metaflow.decorators import StepDecorator
+from metaflow import current
+from .core import AppDeployer, apps
+from .core.perimeters import PerimeterExtractor
+import os
+import hashlib
+
+
+class AppDeployDecorator(StepDecorator):
+
+    """
+    MF Add To Current
+    -----------------
+    apps -> metaflow_extensions.outerbounds.plugins.apps.core.apps
+
+        @@ Returns
+        ----------
+        apps
+            The object carrying the Deployer class to deploy apps.
+    """
+
+    name = "app_deploy"
+    defaults = {}
+
+    package_url = None
+    package_sha = None
+
+    MAX_ENTROPY = 6
+    MAX_NAME_LENGTH = 15 - MAX_ENTROPY - 1  # -1 for the hyphen
+
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        self.logger = logger
+        self.environment = environment
+        self.step = step
+        self.flow_datastore = flow_datastore
+
+    def _resolve_package_url_and_sha(self):
+        return os.environ.get("METAFLOW_CODE_URL", self.package_url), os.environ.get(
+            "METAFLOW_CODE_SHA", self.package_sha
+        )
+
+    def _extract_project_info(self):
+        project = current.get("project_name")
+        branch = current.get("branch_name")
+        is_production = current.get("is_production")
+        return project, branch, is_production
+
+    def _resolve_default_image(self, flow):
+        # TODO : Resolve the default image over here.
+        pass
+
+    def _resolve_default_name_prefix(self, flow, step_name):
+        # TODO: Only tweek MAX_NAME_LENGTH as backend support allows longer names.
+        base_prefix = (flow.name + "-" + step_name).lower()
+        if len(base_prefix) > self.MAX_NAME_LENGTH:
+            base_prefix = "mf-app"
+        return base_prefix
+
+    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,
+    ):
+        perimeter, api_server = PerimeterExtractor.during_metaflow_execution()
+        package_url, package_sha = self._resolve_package_url_and_sha()
+        if package_url is None or package_sha is None:
+            raise MetaflowException(
+                "METAFLOW_CODE_URL or METAFLOW_CODE_SHA is not set. "
+                "Please set METAFLOW_CODE_URL and METAFLOW_CODE_SHA in your environment."
+            )
+        image = os.environ.get("FASTBAKERY_IMAGE", None)
+
+        # TODO [Apps] - This is temporary. Backend will support longer names in the future.
+        default_name = self._resolve_default_name_prefix(flow, step_name)
+        project, branch, is_production = self._extract_project_info()
+        project_info = {}
+        if project is not None:
+            project_info["metaflow/project"] = project
+            project_info["metaflow/branch"] = branch
+            project_info["metaflow/is_production"] = is_production
+
+        default_tags = {
+            "metaflow/flow_name": flow.name,
+            "metaflow/step_name": step_name,
+            "metaflow/run_id": run_id,
+            "metaflow/task_id": task_id,
+            "metaflow/retry_count": retry_count,
+            "metaflow/pathspec": current.pathspec,
+            **project_info,
+        }
+
+        AppDeployer._set_state(
+            perimeter,
+            api_server,
+            code_package_url=package_url,
+            code_package_key=package_sha,
+            name_prefix=default_name,
+            image=image,
+            max_entropy=self.MAX_ENTROPY,
+            default_tags=[{k: str(v)} for k, v in default_tags.items()],
+        )
+        current._update_env(
+            {
+                "apps": apps(),
+            }
+        )
+
+    def task_post_step(
+        self, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        pass
+
+    def runtime_init(self, flow, graph, package, run_id):
+        # Set some more internal state.
+        self.flow = flow
+        self.graph = graph
+        self.package = package
+        self.run_id = run_id
+
+    def runtime_task_created(
+        self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
+    ):
+        # To execute the Kubernetes job, the job container needs to have
+        # access to the code package. We store the package in the datastore
+        # which the pod is able to download as part of it's entrypoint.
+        if not is_cloned:
+            self._save_package_once(self.flow_datastore, self.package)
+
+    @classmethod
+    def _save_package_once(cls, flow_datastore, package):
+        if cls.package_url is None:
+            cls.package_url, cls.package_sha = flow_datastore.save_data(
+                [package.blob], len_hint=1
+            )[0]
+            os.environ["METAFLOW_CODE_URL"] = cls.package_url
+            os.environ["METAFLOW_CODE_SHA"] = cls.package_sha

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

@@ -0,0 +1,10 @@
+from . import app_cli
+from . import config
+from .deployer import AppDeployer, apps
+from .config.typed_configs import (
+    ReplicaConfigDict,
+    ResourceConfigDict,
+    AuthConfigDict,
+    DependencyConfigDict,
+    PackageConfigDict,
+)