ob-metaflow-extensions 1.1.163rc6__py2.py3-none-any.whl → 1.1.164__py2.py3-none-any.whl

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
@@ -49,6 +52,10 @@ def get_boto3_session(role_arn=None, session_vars=None):
 
     token_info = get_token("/generate/aws")
 
+    # Check if the assume_role decorator has set a role ARN via environment variable
+    # This takes precedence over CSPR role
+    decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+
     # Write token to a file. The file name is derived from the user name
     # so it works with multiple users on the same machine.
     #
@@ -69,12 +76,13 @@ def get_boto3_session(role_arn=None, session_vars=None):
     if token_info.get("cspr_role_arn"):
         cspr_role = token_info["cspr_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
+    # If assume_role decorator is used, prioritize it over CSPR role
+    effective_role = decorator_role_arn or cspr_role
+    if effective_role:
+        # If we have either a decorator role or CSPR role, set up AWS config
         # 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.
+        # effective role using the task role credentials.
         import configparser
         from io import StringIO
 
@@ -86,9 +94,9 @@ def get_boto3_session(role_arn=None, session_vars=None):
             "web_identity_token_file": token_file,
         }
 
-        # CSPR role profile (default)
-        aws_config["profile cspr"] = {
-            "role_arn": cspr_role,
+        # Effective role profile (decorator role or CSPR role)
+        aws_config["profile effective"] = {
+            "role_arn": effective_role,
             "source_profile": "task",
         }
 
@@ -104,7 +112,7 @@ def get_boto3_session(role_arn=None, session_vars=None):
             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"
+        os.environ["AWS_PROFILE"] = "effective"
     else:
         os.environ["AWS_WEB_IDENTITY_TOKEN_FILE"] = token_file
         os.environ["AWS_ROLE_ARN"] = token_info["role_arn"]
@@ -116,21 +124,29 @@ def get_boto3_session(role_arn=None, session_vars=None):
     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")
+    if effective_role:
+        if role_arn == decorator_role_arn or role_arn == USE_CSPR_ROLE_ARN_IF_SET:
+            # We have either a decorator role or CSPR role, use the effective profile
+            # The generated AWS config will be used here since we set the
+            # AWS_CONFIG_FILE environment variable above.
+            session = boto3.session.Session(profile_name="effective")
         else:
             session = boto3.session.Session(profile_name="task")
     else:
+        # No decorator role or CSPR role, use default session
         # 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.
+    if (
+        role_arn
+        and role_arn != USE_CSPR_ROLE_ARN_IF_SET
+        and role_arn != decorator_role_arn
+    ):
+        # If the user provided a role_arn that's different from the decorator role,
+        # we assume that role using the current session credentials.
+        # This works for both cases:
+        # 1. No decorator role: Task role -> Secrets role
+        # 2. With decorator role: Decorator role -> Secrets role
         fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
             client_creator=session._session.create_client,
             source_credentials=session._session.get_credentials(),
@@ -146,8 +162,8 @@ def get_boto3_session(role_arn=None, session_vars=None):
     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.
+        # session which would use the effective role (decorator or CSPR) if set,
+        # or the task role otherwise.
         return session
 
 

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,65 @@
+from metaflow.user_configs.config_decorators import (
+    MutableFlow,
+    MutableStep,
+    CustomFlowDecorator,
+)
+from .assume_role import OBP_ASSUME_ROLE_ARN_ENV_VAR
+
+
+class assume_role(CustomFlowDecorator):
+    """
+    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
+    """
+
+    def init(self, *args, **kwargs):
+        self.role_arn = kwargs.get("role_arn", 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::'"
+            )
+
+    def evaluate(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
+
+        # 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:
+            # We'll inject the role assumption by adding an environment decorator
+            # The role will be available through an environment variable
+            step.add_decorator(
+                environment, vars={OBP_ASSUME_ROLE_ARN_ENV_VAR: self.role_arn}
+            )
+
+        # Apply the role assumption setup to all steps in the flow
+        for _, step in mutable_flow.steps:
+            _setup_role_assumption(step)

metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py

@@ -20,8 +20,19 @@ def get_aws_client(
     module, with_error=False, role_arn=None, session_vars=None, client_params=None
 ):
     import metaflow.plugins.aws.aws_client
+    import os
 
     from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+    from metaflow_extensions.outerbounds.plugins.aws.assume_role import (
+        OBP_ASSUME_ROLE_ARN_ENV_VAR,
+    )
+
+    # Check if the assume_role decorator has set a role ARN via environment variable
+    # This takes precedence over CSPR but not over explicitly passed role_arn
+    if role_arn is None:
+        decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+        if decorator_role_arn:
+            role_arn = decorator_role_arn
 
     return metaflow.plugins.aws.aws_client.get_aws_client(
         module,
@@ -42,10 +53,20 @@ def get_aws_client(
 def S3(*args, **kwargs):
     import sys
     import metaflow.plugins.datatools.s3
+    import os
     from metaflow_extensions.outerbounds.plugins import USE_CSPR_ROLE_ARN_IF_SET
+    from metaflow_extensions.outerbounds.plugins.aws.assume_role import (
+        OBP_ASSUME_ROLE_ARN_ENV_VAR,
+    )
 
+    # Check if the assume_role decorator has set a role ARN via environment variable
+    # This takes precedence over CSPR but not over explicitly passed role
     if "role" not in kwargs or kwargs["role"] is None:
-        kwargs["role"] = USE_CSPR_ROLE_ARN_IF_SET
+        decorator_role_arn = os.environ.get(OBP_ASSUME_ROLE_ARN_ENV_VAR)
+        if decorator_role_arn:
+            kwargs["role"] = decorator_role_arn
+        else:
+            kwargs["role"] = USE_CSPR_ROLE_ARN_IF_SET
 
     return metaflow.plugins.datatools.s3.S3(*args, **kwargs)
 
@@ -53,4 +74,4 @@ def S3(*args, **kwargs):
 from .. import profilers
 from ..plugins.snowflake import Snowflake
 from ..plugins.checkpoint_datastores import nebius_checkpoints, coreweave_checkpoints
-from . import ob_internal
+from ..plugins.aws import assume_role

{ob_metaflow_extensions-1.1.163rc6.dist-info → ob_metaflow_extensions-1.1.164.dist-info}/METADATA

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

{ob_metaflow_extensions-1.1.163rc6.dist-info → ob_metaflow_extensions-1.1.164.dist-info}/RECORD

@@ -1,7 +1,7 @@
 metaflow_extensions/outerbounds/__init__.py,sha256=Gb8u06s9ClQsA_vzxmkCzuMnigPy7kKcDnLfb7eB-64,514
 metaflow_extensions/outerbounds/remote_config.py,sha256=pEFJuKDYs98eoB_-ryPjVi9b_c4gpHMdBHE14ltoxIU,4672
 metaflow_extensions/outerbounds/config/__init__.py,sha256=JsQGRuGFz28fQWjUvxUgR8EKBLGRdLUIk_buPLJplJY,1225
-metaflow_extensions/outerbounds/plugins/__init__.py,sha256=GxYKjrMJCGVKoxhfdPAlVF9kYrEb3-xn9fgUTb_H9VY,13255
+metaflow_extensions/outerbounds/plugins/__init__.py,sha256=QERuUPciMFDxUfzZThqArCGVI_lRXSAASUnQVOmlw8E,14058
 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
@@ -9,6 +9,9 @@ metaflow_extensions/outerbounds/plugins/apps/app_utils.py,sha256=sw9whU17lAzlD2K
 metaflow_extensions/outerbounds/plugins/apps/consts.py,sha256=iHsyqbUg9k-rgswCs1Jxf5QZIxR1V-peCDRjgr9kdBM,177
 metaflow_extensions/outerbounds/plugins/apps/deploy_decorator.py,sha256=VkmiMdNYHhNdt-Qm9AVv7aE2LWFsIFEc16YcOYjwF6Q,8568
 metaflow_extensions/outerbounds/plugins/apps/supervisord_utils.py,sha256=GQoN2gyPClcpR9cLldJmbCfqXnoAHxp8xUnY7vzaYtY,9026
+metaflow_extensions/outerbounds/plugins/aws/__init__.py,sha256=VBGdjNKeFLXGZuqh4jVk8cFtO1AWof73a6k_cnbAOYA,145
+metaflow_extensions/outerbounds/plugins/aws/assume_role.py,sha256=mBewNlnSYsR2rFXFkX-DUH6ku01h2yOcMcLHoCL7eyI,161
+metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py,sha256=b3te8-RJxsj5cqXbQTFxxeIfnqTERAFInzLgXLVLyYA,2391
 metaflow_extensions/outerbounds/plugins/card_utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/card_utilities/async_cards.py,sha256=6rQhtZXK5DenXPfCRS1ul0jvLJYd48jrJAlnodID21w,4347
 metaflow_extensions/outerbounds/plugins/card_utilities/extra_components.py,sha256=D8rgCaSc7PuLD0MHJjqsjN0g0PQMN1H-ySOJqi5uIOE,19111
@@ -17,7 +20,6 @@ metaflow_extensions/outerbounds/plugins/checkpoint_datastores/__init__.py,sha256
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.py,sha256=_WzoOROFjoFa8TzsMNFp-r_1Zz7NUp-5ljn_kKlczXA,4534
 metaflow_extensions/outerbounds/plugins/checkpoint_datastores/nebius.py,sha256=zgqDLFewCeF5jqh-hUNKmC_OAjld09ln0bb8Lkeqapc,4659
 metaflow_extensions/outerbounds/plugins/fast_bakery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py,sha256=ShE5omFBr83wkvEhL_ptRFvDNMs6wefg4BjaafQjTcM,3602
 metaflow_extensions/outerbounds/plugins/fast_bakery/docker_environment.py,sha256=Tl520HdBteg-aDOM7mnnJJpdDCZc49BmFFmLUc_vTi8,15018
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery.py,sha256=PE81ZB54OAMXkMGSB7JqgvgMg7N9kvoVclrWL-6jc2U,5626
 metaflow_extensions/outerbounds/plugins/fast_bakery/fast_bakery_cli.py,sha256=kqFyu2bJSnc9_9aYfBpz5xK6L6luWFZK_NMuh8f1eVk,1494
@@ -68,15 +70,14 @@ metaflow_extensions/outerbounds/plugins/torchtune/__init__.py,sha256=TOXNeyhcgd8
 metaflow_extensions/outerbounds/profilers/__init__.py,sha256=wa_jhnCBr82TBxoS0e8b6_6sLyZX0fdHicuGJZNTqKw,29
 metaflow_extensions/outerbounds/profilers/gpu.py,sha256=3Er8uKQzfm_082uadg4yn_D4Y-iSCgzUfFmguYxZsz4,27485
 metaflow_extensions/outerbounds/toplevel/__init__.py,sha256=qWUJSv_r5hXJ7jV_On4nEasKIfUCm6_UjkjXWA_A1Ts,90
-metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=505fLAY4NkBCGtjpXlE0RgaKVpJ0jOUq1-Fq-EKzPew,2035
-metaflow_extensions/outerbounds/toplevel/ob_internal.py,sha256=53xM6d_UYT3uGFFA59UzxN23H5QMO5_F39pALpmGy04,51
+metaflow_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.py,sha256=FS0ZKQJsJKlw9PgtLqVV4kH9o7_qwYcIVtyu2Kqwa_U,2973
 metaflow_extensions/outerbounds/toplevel/plugins/azure/__init__.py,sha256=WUuhz2YQfI4fz7nIcipwwWq781eaoHEk7n4GAn1npDg,63
 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/ollama/__init__.py,sha256=GRSz2zwqkvlmFS6bcfYD_CX6CMko9DHQokMaH1iBshA,47
 metaflow_extensions/outerbounds/toplevel/plugins/snowflake/__init__.py,sha256=LptpH-ziXHrednMYUjIaosS1SXD3sOtF_9_eRqd8SJw,50
 metaflow_extensions/outerbounds/toplevel/plugins/torchtune/__init__.py,sha256=uTVkdSk3xZ7hEKYfdlyVteWj5KeDwaM1hU9WT-_YKfI,50
-ob_metaflow_extensions-1.1.163rc6.dist-info/METADATA,sha256=IGm75le9ACGIrfLPPPMKmqKsDrhkoAex5xZ7KDikMr4,524
-ob_metaflow_extensions-1.1.163rc6.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.1.163rc6.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.1.163rc6.dist-info/RECORD,,
+ob_metaflow_extensions-1.1.164.dist-info/METADATA,sha256=SEMbmxC4QBlLcwZQ8Wr9YKxoi6dl6KJoHRtau7Zsx4k,521
+ob_metaflow_extensions-1.1.164.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.1.164.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.1.164.dist-info/RECORD,,

metaflow_extensions/outerbounds/plugins/fast_bakery/baker.py

@@ -1,110 +0,0 @@
-import threading
-import time
-import sys
-from typing import Dict, Optional, Any, Callable
-from functools import partial
-from metaflow.exception import MetaflowException
-from metaflow.metaflow_config import FAST_BAKERY_URL
-
-from .fast_bakery import FastBakery, FastBakeryApiResponse, FastBakeryException
-from .docker_environment import cache_request
-
-BAKERY_METAFILE = ".imagebakery-cache"
-
-
-class BakerException(MetaflowException):
-    headline = "Ran into an error while baking image"
-
-    def __init__(self, msg):
-        super(BakerException, self).__init__(msg)
-
-
-def bake_image(
-    cache_file_path: str,
-    ref: Optional[str] = None,
-    python: Optional[str] = None,
-    pypi_packages: Optional[Dict[str, str]] = None,
-    conda_packages: Optional[Dict[str, str]] = None,
-    base_image: Optional[str] = None,
-    logger: Optional[Callable[[str], Any]] = None,
-) -> FastBakeryApiResponse:
-    """
-    Bakes a Docker image with the specified dependencies.
-
-    Args:
-        cache_file_path: Path to the cache file
-        ref: Reference identifier for this bake (for logging purposes)
-        python: Python version to use
-        pypi_packages: Dictionary of PyPI packages and versions
-        conda_packages: Dictionary of Conda packages and versions
-        base_image: Base Docker image to use
-        logger: Optional logger function to output progress
-
-    Returns:
-        FastBakeryApiResponse: The response from the bakery service
-
-    Raises:
-        BakerException: If the baking process fails
-    """
-    # Default logger if none provided
-    if logger is None:
-        logger = partial(print, file=sys.stderr)
-
-    # Thread lock for logging
-    logger_lock = threading.Lock()
-    images_baked = 0
-
-    @cache_request(cache_file_path)
-    def _cached_bake(
-        ref=None,
-        python=None,
-        pypi_packages=None,
-        conda_packages=None,
-        base_image=None,
-    ):
-        try:
-            bakery = FastBakery(url=FAST_BAKERY_URL)
-            bakery._reset_payload()
-            bakery.python_version(python)
-            bakery.pypi_packages(pypi_packages)
-            bakery.conda_packages(conda_packages)
-            bakery.base_image(base_image)
-            # bakery.ignore_cache()
-
-            with logger_lock:
-                logger(f"🍳 Baking [{ref}] ...")
-                logger(f"     🐍 Python: {python}")
-
-                if pypi_packages:
-                    logger(f"     📦 PyPI packages:")
-                    for package, version in pypi_packages.items():
-                        logger(f"        🔧 {package}: {version}")
-
-                if conda_packages:
-                    logger(f"     📦 Conda packages:")
-                    for package, version in conda_packages.items():
-                        logger(f"        🔧 {package}: {version}")
-
-                logger(f"     🏗️  Base image: {base_image}")
-
-            start_time = time.time()
-            res = bakery.bake()
-            # TODO: Get actual bake time from bakery
-            bake_time = time.time() - start_time
-
-            with logger_lock:
-                logger(f"🏁 Baked [{ref}] in {bake_time:.2f} seconds!")
-            nonlocal images_baked
-            images_baked += 1
-            return res
-        except FastBakeryException as ex:
-            raise BakerException(f"Bake [{ref}] failed: {str(ex)}")
-
-    # Call the cached bake function with the provided parameters
-    return _cached_bake(
-        ref=ref,
-        python=python,
-        pypi_packages=pypi_packages,
-        conda_packages=conda_packages,
-        base_image=base_image,
-    )

metaflow_extensions/outerbounds/toplevel/ob_internal.py

@@ -1 +0,0 @@
-from ..plugins.fast_bakery.baker import bake_image