ob-metaflow-extensions 1.4.35__py2.py3-none-any.whl → 1.4.39__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py

@@ -25,10 +25,29 @@ class assume_role(FlowMutator):
         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(
@@ -40,6 +59,13 @@ class assume_role(FlowMutator):
                 "`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.
@@ -49,6 +75,18 @@ class assume_role(FlowMutator):
         # 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}
@@ -73,6 +111,8 @@ class assume_role(FlowMutator):
         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
-        for _, step in mutable_flow.steps:
-            _setup_role_assumption(step)
+        # 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)

metaflow_extensions/outerbounds/plugins/snowflake/snowflake.py

@@ -83,22 +83,32 @@ def get_snowflake_token(user: str = "", role: str = "", integration: str = "") -
     return response.json()["token"]
 
 
-def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
+def get_oauth_connection_params(
+    user: str = "", role: str = "", integration: str = "", **kwargs
+) -> Dict:
     """
-    Connect to snowflake using the token minted by Outerbounds
+    Get OAuth connection parameters for Snowflake authentication using Outerbounds integration.
+
+    This is a helper function that returns connection parameters dict that can be used
+    with both snowflake-connector-python and snowflake-snowpark-python.
+
     user: str
         The user name used to authenticate with snowflake
     role: str
-        The role to request when connect with snowflake
+        The role to request when connecting with snowflake
     integration: str
-        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists in the current perimeter. If integration is not set and more than one exists in the current perimeter, then we raise an exception.
+        The name of the snowflake integration to use. If not set, an existing integration
+        will be used provided that only one exists in the current perimeter.
     kwargs: dict
-        Additional arguments to pass to the python snowflake connector
+        Additional arguments to include in the connection parameters
+
+    Returns:
+        Dict with connection parameters including OAuth token
     """
     # ensure password is not set
     if "password" in kwargs:
         raise OuterboundsSnowflakeConnectorException(
-            "Password should not be set when using Outerbounds snowflake connector."
+            "Password should not be set when using Outerbounds OAuth authentication."
         )
 
     provisioner = SnowflakeIntegrationProvisioner(integration)
@@ -137,11 +147,31 @@ def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
     kwargs["role"] = role
     kwargs["user"] = user
 
+    return kwargs
+
+
+def connect(user: str = "", role: str = "", integration: str = "", **kwargs):
+    """
+    Connect to snowflake using the token minted by Outerbounds
+    user: str
+        The user name used to authenticate with snowflake
+    role: str
+        The role to request when connect with snowflake
+    integration: str
+        The name of the snowflake integration to use. If not set, an existing integration will be used provided that only one exists in the current perimeter. If integration is not set and more than one exists in the current perimeter, then we raise an exception.
+    kwargs: dict
+        Additional arguments to pass to the python snowflake connector
+    """
+    # Get OAuth connection params using the helper
+    connection_params = get_oauth_connection_params(
+        user=user, role=role, integration=integration, **kwargs
+    )
+
     # connect to snowflake
     try:
         from snowflake.connector import connect
 
-        cn = connect(**kwargs)
+        cn = connect(**connection_params)
         return cn
     except ImportError as ie:
         raise OuterboundsSnowflakeConnectorException(

metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py

@@ -1,9 +1,11 @@
 import os
+import re
 import shlex
 import atexit
 import json
 import math
 import time
+import hashlib
 
 from metaflow import util
 
@@ -57,21 +59,29 @@ class Snowpark(object):
         atexit.register(lambda: self.job.kill() if hasattr(self, "job") else None)
 
     def _job_name(self, user, flow_name, run_id, step_name, task_id, retry_count):
-        return "{user}-{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}".format(
-            user=user,
-            flow_name=flow_name,
-            run_id=str(run_id) if run_id is not None else "",
-            step_name=step_name,
-            task_id=str(task_id) if task_id is not None else "",
-            retry_count=str(retry_count) if retry_count is not None else "",
+        unique_str = (
+            "{user}-{flow_name}-{run_id}-{step_name}-{task_id}-{retry_count}".format(
+                user=user,
+                flow_name=flow_name,
+                run_id=str(run_id) if run_id is not None else "",
+                step_name=step_name,
+                task_id=str(task_id) if task_id is not None else "",
+                retry_count=str(retry_count) if retry_count is not None else "",
+            )
         )
+        unique_hash = hashlib.md5(unique_str.encode("utf-8")).hexdigest()[:8]
+        raw_prefix = f"{flow_name}-{step_name}"
+        safe_prefix = re.sub(r"[^a-z0-9]", "-", raw_prefix.lower())
+        safe_prefix = safe_prefix[:54]
+        safe_prefix = safe_prefix.lstrip("-")
+        return f"{safe_prefix}-{unique_hash}"
 
     def _command(self, environment, code_package_url, step_name, step_cmds, task_spec):
         mflog_expr = export_mflog_env_vars(
             datastore_type=self.datastore.TYPE,
             stdout_path=STDOUT_PATH,
             stderr_path=STDERR_PATH,
-            **task_spec
+            **task_spec,
         )
         init_cmds = environment.get_package_commands(
             code_package_url, self.datastore.TYPE

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py

@@ -66,6 +66,10 @@ def snowpark():
     "--schema",
     help="Schema for Snowpark Container Services.",
 )
+@click.option(
+    "--integration",
+    help="Outerbounds OAuth integration name for Snowpark Container Services. When set, uses OAuth authentication instead of password.",
+)
 @click.option(
     "--image",
     help="Docker image requirement for Snowpark Container Services. In name:version format.",
@@ -119,6 +123,7 @@ def step(
     database=None,
     warehouse=None,
     schema=None,
+    integration=None,
     image=None,
     stage=None,
     compute_pool=None,
@@ -235,6 +240,7 @@ def step(
                 "database": database,
                 "warehouse": warehouse,
                 "schema": schema,
+                "integration": integration,
             },
         )
         with ctx.obj.monitor.measure("metaflow.snowpark.launch_job"):

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py

@@ -10,14 +10,15 @@ from metaflow.exception import MetaflowException
 class SnowparkClient(object):
     def __init__(
         self,
-        account: str,
-        user: str,
-        password: str,
-        role: str,
-        database: str,
-        warehouse: str,
-        schema: str,
+        account: str = None,
+        user: str = None,
+        password: str = None,
+        role: str = None,
+        database: str = None,
+        warehouse: str = None,
+        schema: str = None,
         autocommit: bool = True,
+        integration: str = None,
     ):
         try:
             from snowflake.core import Root
@@ -37,15 +38,38 @@ class SnowparkClient(object):
                 % sys.executable
             )
 
+        if integration:
+            # Use OAuth authentication via Outerbounds integration
+            from metaflow_extensions.outerbounds.plugins.snowflake.snowflake import (
+                get_oauth_connection_params,
+            )
+
+            self.connection_parameters = get_oauth_connection_params(
+                user=user or "",
+                role=role or "",
+                integration=integration,
+                schema=schema or "",
+                account=account,
+                warehouse=warehouse,
+                database=database,
+            )
+            self.connection_parameters["autocommit"] = autocommit
+        else:
+            # Password-based authentication
+            self.connection_parameters = {
+                "account": account,
+                "user": user,
+                "password": password,
+                "role": role,
+                "warehouse": warehouse,
+                "database": database,
+                "schema": schema,
+                "autocommit": autocommit,
+            }
+
+        # Remove None values from connection parameters
         self.connection_parameters = {
-            "account": account,
-            "user": user,
-            "password": password,
-            "role": role,
-            "warehouse": warehouse,
-            "database": database,
-            "schema": schema,
-            "autocommit": autocommit,
+            k: v for k, v in self.connection_parameters.items() if v is not None
         }
 
         try:

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py

@@ -71,6 +71,7 @@ class SnowparkDecorator(StepDecorator):
         "cpu": None,
         "gpu": None,
         "memory": None,
+        "integration": None,  # Outerbounds OAuth integration name
     }
 
     package_url = None
@@ -80,12 +81,11 @@ class SnowparkDecorator(StepDecorator):
     def __init__(self, attributes=None, statically_defined=False):
         super(SnowparkDecorator, self).__init__(attributes, statically_defined)
 
+        # Set defaults from config (user can override via decorator or integration)
         if not self.attributes["account"]:
             self.attributes["account"] = SNOWPARK_ACCOUNT
         if not self.attributes["user"]:
             self.attributes["user"] = SNOWPARK_USER
-        if not self.attributes["password"]:
-            self.attributes["password"] = SNOWPARK_PASSWORD
         if not self.attributes["role"]:
             self.attributes["role"] = SNOWPARK_ROLE
         if not self.attributes["database"]:
@@ -94,6 +94,9 @@ class SnowparkDecorator(StepDecorator):
             self.attributes["warehouse"] = SNOWPARK_WAREHOUSE
         if not self.attributes["schema"]:
             self.attributes["schema"] = SNOWPARK_SCHEMA
+        # Only use password from config if not using integration (OAuth)
+        if not self.attributes["integration"] and not self.attributes["password"]:
+            self.attributes["password"] = SNOWPARK_PASSWORD
 
         # If no docker image is explicitly specified, impute a default image.
         if not self.attributes["image"]:

metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py

@@ -12,9 +12,9 @@ from .snowpark_exceptions import SnowparkException
 mapping = str.maketrans("0123456789", "abcdefghij")
 
 
-# keep only alpha numeric characters and underscores..
+# keep only alpha numeric characters and dashes..
 def sanitize_name(job_name: str):
-    return "".join(char for char in job_name if char.isalnum() or char == "_")
+    return "".join(char for char in job_name if char.isalnum() or char == "-")
 
 
 # this is not a decorator since the exception imports need to be inside

{ob_metaflow_extensions-1.4.35.dist-info → ob_metaflow_extensions-1.4.39.dist-info}/METADATA

@@ -1,13 +1,13 @@
 Metadata-Version: 2.1
 Name: ob-metaflow-extensions
-Version: 1.4.35
+Version: 1.4.39
 Summary: Outerbounds Platform Extensions for Metaflow
 Author: Outerbounds, Inc.
 License: Commercial
 Description-Content-Type: text/markdown
 Requires-Dist: boto3
 Requires-Dist: kubernetes
-Requires-Dist: ob-metaflow (==2.19.12.3)
+Requires-Dist: ob-metaflow (==2.19.15.1)
 
 # Outerbounds platform package
 

{ob_metaflow_extensions-1.4.35.dist-info → ob_metaflow_extensions-1.4.39.dist-info}/RECORD

@@ -41,7 +41,7 @@ metaflow_extensions/outerbounds/plugins/apps/core/config/unified_config.py,sha25
 metaflow_extensions/outerbounds/plugins/apps/core/experimental/__init__.py,sha256=_OT0RsbQdYerT-fjG_vcu85ava6sHbCsGBpLcGcrlX8,3043
 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=MrExIdwnQ4ZVKbrWqEQAIuZO4lWEy1rE0ihGUx7GMJA,2976
+metaflow_extensions/outerbounds/plugins/aws/assume_role_decorator.py,sha256=cZg2zUcyvehpjE-a12HTP6vpwPG2MAJuK5KSywBk8xs,4638
 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
@@ -98,14 +98,14 @@ metaflow_extensions/outerbounds/plugins/s3_proxy/s3_proxy_manager.py,sha256=Grl7
 metaflow_extensions/outerbounds/plugins/secrets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 metaflow_extensions/outerbounds/plugins/secrets/secrets.py,sha256=3s98hO_twKkM22tKyDdcUjGQNfYpSXW_jLKISV9ju_U,8433
 metaflow_extensions/outerbounds/plugins/snowflake/__init__.py,sha256=RG4ixt3jwqcK1_tt0QxLcUbNmf7wWAMnZhBx-ZMGgLk,114
-metaflow_extensions/outerbounds/plugins/snowflake/snowflake.py,sha256=zoWSHM4CJSfUmJSP-_i4zREWyQOW4USBlgjhQnEhlTE,13669
+metaflow_extensions/outerbounds/plugins/snowflake/snowflake.py,sha256=_ULhLrKZ4rOVcJFb_EmMGi7GFDN6Vtawz8o86mLBMq8,14756
 metaflow_extensions/outerbounds/plugins/snowpark/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py,sha256=0R8aFN9MpgWraqiaI6ZF82YpLdFJ1f-3z_-BPRpZfxM,10674
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py,sha256=ezJ2Jr8JJ48SvTH0ET7pRUGaBuq3XTSfNxJVgXhZsnY,8756
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py,sha256=JWA0n3imeMP3HoDw9_9Ge9UPMAwY_TnPKaETHDN8qjA,4460
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py,sha256=SMND-1BJoqMkYFfFa3B_OT6Xu6HH07degq3qXWOrJ2w,10167
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark.py,sha256=tVpSOizbVjoXp-iMWuzJgE41hUstMrPC8O_eCKY9Mg4,11083
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark_cli.py,sha256=aesXYZGoEF8UtAIvCDJNLTSc9pM62KuOqd9Za95HjZ4,8997
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark_client.py,sha256=6U1sdXx0pdhNdkGba8DaE7GabQTc5fyJI3LpPsPQt0s,5463
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark_decorator.py,sha256=EN-XQafl14uQ3oVGgrkNII5PCYq2lWOt3t2d-JGPeBk,10430
 metaflow_extensions/outerbounds/plugins/snowpark/snowpark_exceptions.py,sha256=FTfYlJu-sn9DkPOs2R1V1ChWb1vZthOgeq0BZdT1ucY,296
-metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=k0OLbsExQOcLbU3_qBd-p2-HCgzcWqXUHxj6L5b6FOk,7075
+metaflow_extensions/outerbounds/plugins/snowpark/snowpark_job.py,sha256=KcbAYZdInAlS10IAZWdqtWAJXY4I6N2ebCTXzAmT72o,7070
 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/plugins/torchtune/__init__.py,sha256=Psj2ybj_E1qp5KK2inon9e4ZecaRxnPtW3ngcirbO2g,6094
@@ -128,7 +128,7 @@ metaflow_extensions/outerbounds/toplevel/plugins/optuna/__init__.py,sha256=6D1wL
 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
 metaflow_extensions/outerbounds/toplevel/plugins/vllm/__init__.py,sha256=ekcgD3KVydf-a0xMI60P4uy6ePkSEoFHiGnDq1JM940,45
-ob_metaflow_extensions-1.4.35.dist-info/METADATA,sha256=DVnCm-4Hom7YA49NHtsBa1P04fIuv9rECbUpVhRY0Nc,520
-ob_metaflow_extensions-1.4.35.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-ob_metaflow_extensions-1.4.35.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
-ob_metaflow_extensions-1.4.35.dist-info/RECORD,,
+ob_metaflow_extensions-1.4.39.dist-info/METADATA,sha256=Zfm2ymOc5lvyeuNPZUNpmCHylRjYHRKGQVhyL_YduBw,520
+ob_metaflow_extensions-1.4.39.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+ob_metaflow_extensions-1.4.39.dist-info/top_level.txt,sha256=NwG0ukwjygtanDETyp_BUdtYtqIA_lOjzFFh1TsnxvI,20
+ob_metaflow_extensions-1.4.39.dist-info/RECORD,,