ob-metaflow 2.11.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/aws/step_functions/step_functions.py

@@ -18,6 +18,7 @@ from metaflow.metaflow_config import (
     SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH,
 )
 from metaflow.parameters import deploy_time_eval
+from metaflow.user_configs.config_options import ConfigInput
 from metaflow.util import dict_to_cli_options, to_pascalcase
 
 from ..batch.batch import Batch
@@ -39,6 +40,7 @@ class StepFunctions(object):
         name,
         graph,
         flow,
+        code_package_metadata,
         code_package_sha,
         code_package_url,
         production_token,
@@ -48,16 +50,19 @@ class StepFunctions(object):
         event_logger,
         monitor,
         tags=None,
+        aws_batch_tags=None,
         namespace=None,
         username=None,
         max_workers=None,
         workflow_timeout=None,
         is_project=False,
         use_distributed_map=False,
+        compress_state_machine=False,
     ):
         self.name = name
         self.graph = graph
         self.flow = flow
+        self.code_package_metadata = code_package_metadata
         self.code_package_sha = code_package_sha
         self.code_package_url = code_package_url
         self.production_token = production_token
@@ -67,14 +72,19 @@ class StepFunctions(object):
         self.event_logger = event_logger
         self.monitor = monitor
         self.tags = tags
+        self.aws_batch_tags = aws_batch_tags or {}
         self.namespace = namespace
         self.username = username
         self.max_workers = max_workers
         self.workflow_timeout = workflow_timeout
+        self.config_parameters = self._process_config_parameters()
 
         # https://aws.amazon.com/blogs/aws/step-functions-distributed-map-a-serverless-solution-for-large-scale-parallel-data-processing/
         self.use_distributed_map = use_distributed_map
 
+        # S3 command upload configuration
+        self.compress_state_machine = compress_state_machine
+
         self._client = StepFunctionsClient()
         self._workflow = self._compile()
         self._cron = self._cron()
@@ -190,6 +200,7 @@ class StepFunctions(object):
                 "on AWS Step Functions. Please "
                 "deploy your flow first." % name
             )
+
         # Dump parameters into `Parameters` input field.
         input = json.dumps({"Parameters": json.dumps(parameters)})
         # AWS Step Functions limits input to be 32KiB, but AWS Batch
@@ -234,7 +245,7 @@ class StepFunctions(object):
                 return parameters.get("metaflow.owner"), parameters.get(
                     "metaflow.production_token"
                 )
-            except KeyError as e:
+            except KeyError:
                 raise StepFunctionsException(
                     "An existing non-metaflow "
                     "workflow with the same name as "
@@ -299,6 +310,12 @@ class StepFunctions(object):
                 "to AWS Step Functions is not supported currently."
             )
 
+        if self.flow._flow_decorators.get("exit_hook"):
+            raise StepFunctionsException(
+                "Deploying flows with the @exit_hook decorator "
+                "to AWS Step Functions is not currently supported."
+            )
+
         # Visit every node of the flow and recursively build the state machine.
         def _visit(node, workflow, exit_node=None):
             if node.parallel_foreach:
@@ -307,6 +324,12 @@ class StepFunctions(object):
                     "to AWS Step Functions is not supported currently."
                 )
 
+            if node.type == "split-switch":
+                raise StepFunctionsException(
+                    "Deploying flows with switch statement "
+                    "to AWS Step Functions is not supported currently."
+                )
+
             # Assign an AWS Batch job to the AWS Step Functions state
             # and pass the intermediate state by exposing `JobId` and
             # `Parameters` to the child job(s) as outputs. `Index` and
@@ -485,6 +508,10 @@ class StepFunctions(object):
                     "case-insensitive." % param.name
                 )
             seen.add(norm)
+            # NOTE: We skip config parameters as these do not have dynamic values,
+            # and need to be treated differently.
+            if param.IS_CONFIG_PARAMETER:
+                continue
 
             is_required = param.kwargs.get("required", False)
             # Throw an exception if a schedule is set for a flow with required
@@ -501,6 +528,27 @@ class StepFunctions(object):
             parameters.append(dict(name=param.name, value=value))
         return parameters
 
+    def _process_config_parameters(self):
+        parameters = []
+        seen = set()
+        for var, param in self.flow._get_parameters():
+            if not param.IS_CONFIG_PARAMETER:
+                continue
+            # Throw an exception if the parameter is specified twice.
+            norm = param.name.lower()
+            if norm in seen:
+                raise MetaflowException(
+                    "Parameter *%s* is specified twice. "
+                    "Note that parameter names are "
+                    "case-insensitive." % param.name
+                )
+            seen.add(norm)
+
+            parameters.append(
+                dict(name=param.name, kv_name=ConfigInput.make_key_name(param.name))
+            )
+        return parameters
+
     def _batch(self, node):
         attrs = {
             # metaflow.user is only used for setting the AWS Job Name.
@@ -664,9 +712,9 @@ class StepFunctions(object):
                 # input to those descendent tasks. We set and propagate the
                 # task ids pointing to split_parents through every state.
                 if any(self.graph[n].type == "foreach" for n in node.in_funcs):
-                    attrs[
-                        "split_parent_task_id_%s.$" % node.split_parents[-1]
-                    ] = "$.SplitParentTaskId"
+                    attrs["split_parent_task_id_%s.$" % node.split_parents[-1]] = (
+                        "$.SplitParentTaskId"
+                    )
                     for parent in node.split_parents[:-1]:
                         if self.graph[parent].type == "foreach":
                             attrs["split_parent_task_id_%s.$" % parent] = (
@@ -747,6 +795,11 @@ class StepFunctions(object):
         metaflow_version["production_token"] = self.production_token
         env["METAFLOW_VERSION"] = json.dumps(metaflow_version)
 
+        # map config values
+        cfg_env = {param["name"]: param["kv_name"] for param in self.config_parameters}
+        if cfg_env:
+            env["METAFLOW_FLOW_CONFIG_VALUE"] = json.dumps(cfg_env)
+
         # Set AWS DynamoDb Table Name for state tracking for for-eaches.
         # There are three instances when metaflow runtime directly interacts
         # with AWS DynamoDB.
@@ -806,15 +859,17 @@ class StepFunctions(object):
             # AWS_BATCH_JOB_ATTEMPT as the job counter.
             "retry_count": "$((AWS_BATCH_JOB_ATTEMPT-1))",
         }
-
+        # merge batch tags supplied through step-fuctions CLI and ones defined in decorator
+        batch_tags = {**self.aws_batch_tags, **resources["aws_batch_tags"]}
         return (
-            Batch(self.metadata, self.environment)
+            Batch(self.metadata, self.environment, self.flow_datastore)
             .create_job(
                 step_name=node.name,
                 step_cli=self._step_cli(
                     node, input_paths, self.code_package_url, user_code_retries
                 ),
                 task_spec=task_spec,
+                code_package_metadata=self.code_package_metadata,
                 code_package_sha=self.code_package_sha,
                 code_package_url=self.code_package_url,
                 code_package_ds=self.flow_datastore.TYPE,
@@ -831,6 +886,7 @@ class StepFunctions(object):
                 swappiness=resources["swappiness"],
                 efa=resources["efa"],
                 use_tmpfs=resources["use_tmpfs"],
+                aws_batch_tags=batch_tags,
                 tmpfs_tempdir=resources["tmpfs_tempdir"],
                 tmpfs_size=resources["tmpfs_size"],
                 tmpfs_path=resources["tmpfs_path"],
@@ -842,6 +898,7 @@ class StepFunctions(object):
                 ephemeral_storage=resources["ephemeral_storage"],
                 log_driver=resources["log_driver"],
                 log_options=resources["log_options"],
+                offload_command_to_s3=self.compress_state_machine,
             )
             .attempts(total_retries + 1)
         )
@@ -875,13 +932,13 @@ class StepFunctions(object):
             "with": [
                 decorator.make_decorator_spec()
                 for decorator in node.decorators
-                if not decorator.statically_defined
+                if not decorator.statically_defined and decorator.inserted_by is None
             ]
         }
         # FlowDecorators can define their own top-level options. They are
         # responsible for adding their own top-level options and values through
         # the get_top_level_options() hook. See similar logic in runtime.py.
-        for deco in flow_decorators():
+        for deco in flow_decorators(self.flow):
             top_opts_dict.update(deco.get_top_level_options())
 
         top_opts = list(dict_to_cli_options(top_opts_dict))

metaflow/plugins/aws/step_functions/step_functions_cli.py

@@ -7,8 +7,10 @@ from metaflow import JSONType, current, decorators, parameters
 from metaflow._vendor import click
 from metaflow.exception import MetaflowException, MetaflowInternalError
 from metaflow.metaflow_config import (
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
     SERVICE_VERSION_CHECK,
     SFN_STATE_MACHINE_PREFIX,
+    SFN_COMPRESS_STATE_MACHINE,
     UI_URL,
 )
 from metaflow.package import MetaflowPackage
@@ -18,6 +20,8 @@ from metaflow.util import get_username, to_bytes, to_unicode, version_parse
 
 from .production_token import load_token, new_token, store_token
 from .step_functions import StepFunctions
+from metaflow.tagging_util import validate_tags
+from ..aws_utils import validate_aws_tag
 
 VALID_NAME = re.compile(r"[^a-zA-Z0-9_\-\.]")
 
@@ -96,6 +100,13 @@ def step_functions(obj, name=None):
     "with the given tag. You can specify this option multiple "
     "times to attach multiple tags.",
 )
+@click.option(
+    "--aws-batch-tag",
+    "aws_batch_tags",
+    multiple=True,
+    default=None,
+    help="AWS Batch tags.",
+)
 @click.option(
     "--namespace",
     "user_namespace",
@@ -130,10 +141,25 @@ def step_functions(obj, name=None):
     help="Use AWS Step Functions Distributed Map instead of Inline Map for "
     "defining foreach tasks in Amazon State Language.",
 )
+@click.option(
+    "--compress-state-machine/--no-compress-state-machine",
+    is_flag=True,
+    default=SFN_COMPRESS_STATE_MACHINE,
+    help="Compress AWS Step Functions state machine to fit within the 8K limit.",
+)
+@click.option(
+    "--deployer-attribute-file",
+    default=None,
+    show_default=True,
+    type=str,
+    help="Write the workflow name to the file specified. Used internally for Metaflow's Deployer API.",
+    hidden=True,
+)
 @click.pass_obj
 def create(
     obj,
     tags=None,
+    aws_batch_tags=None,
     user_namespace=None,
     only_json=False,
     authorize=None,
@@ -143,9 +169,29 @@ def create(
     workflow_timeout=None,
     log_execution_history=False,
     use_distributed_map=False,
+    compress_state_machine=False,
+    deployer_attribute_file=None,
 ):
+    for node in obj.graph:
+        if any([d.name == "slurm" for d in node.decorators]):
+            raise MetaflowException(
+                "Step *%s* is marked for execution on Slurm with AWS Step Functions which isn't currently supported."
+                % node.name
+            )
+
     validate_tags(tags)
 
+    if deployer_attribute_file:
+        with open(deployer_attribute_file, "w") as f:
+            json.dump(
+                {
+                    "name": obj.state_machine_name,
+                    "flow_name": obj.flow.name,
+                    "metadata": obj.metadata.metadata_str(),
+                },
+                f,
+            )
+
     obj.echo(
         "Deploying *%s* to AWS Step Functions..." % obj.state_machine_name, bold=True
     )
@@ -168,11 +214,13 @@ def create(
         token,
         obj.state_machine_name,
         tags,
+        aws_batch_tags,
         user_namespace,
         max_workers,
         workflow_timeout,
         obj.is_project,
         use_distributed_map,
+        compress_state_machine,
     )
 
     if only_json:
@@ -231,8 +279,10 @@ def check_metadata_service_version(obj):
 
 
 def resolve_state_machine_name(obj, name):
-    def attach_prefix(name):
-        if SFN_STATE_MACHINE_PREFIX is not None:
+    def attach_prefix(name: str):
+        if SFN_STATE_MACHINE_PREFIX is not None and (
+            not name.startswith(SFN_STATE_MACHINE_PREFIX)
+        ):
             return SFN_STATE_MACHINE_PREFIX + "_" + name
         return name
 
@@ -285,32 +335,55 @@ def make_flow(
     token,
     name,
     tags,
+    aws_batch_tags,
     namespace,
     max_workers,
     workflow_timeout,
     is_project,
     use_distributed_map,
+    compress_state_machine=False,
 ):
     if obj.flow_datastore.TYPE != "s3":
         raise MetaflowException("AWS Step Functions requires --datastore=s3.")
 
     # Attach AWS Batch decorator to the flow
     decorators._attach_decorators(obj.flow, [BatchDecorator.name])
+    decorators._init(obj.flow)
     decorators._init_step_decorators(
         obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger
     )
+    obj.graph = obj.flow._graph
 
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=obj.flow_datastore if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE else None,
     )
-    package_url, package_sha = obj.flow_datastore.save_data(
-        [obj.package.blob], len_hint=1
-    )[0]
+    # This blocks until the package is created
+    if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE:
+        package_url = obj.package.package_url()
+        package_sha = obj.package.package_sha()
+    else:
+        package_url, package_sha = obj.flow_datastore.save_data(
+            [obj.package.blob], len_hint=1
+        )[0]
+
+    if aws_batch_tags is not None:
+        batch_tags = {}
+        for item in list(aws_batch_tags):
+            key, value = item.split("=")
+            # These are fresh AWS tags provided by the user through the CLI,
+            # so we must validate them.
+            validate_aws_tag(key, value)
+            batch_tags[key] = value
 
     return StepFunctions(
         name,
         obj.graph,
         obj.flow,
+        obj.package.package_metadata,
         package_sha,
         package_url,
         token,
@@ -320,12 +393,14 @@ def make_flow(
         obj.event_logger,
         obj.monitor,
         tags=tags,
+        aws_batch_tags=batch_tags,
         namespace=namespace,
         max_workers=max_workers,
         username=get_username(),
         workflow_timeout=workflow_timeout,
         is_project=is_project,
         use_distributed_map=use_distributed_map,
+        compress_state_machine=compress_state_machine,
     )
 
 
@@ -440,8 +515,16 @@ def resolve_token(
     type=str,
     help="Write the ID of this run to the file specified.",
 )
+@click.option(
+    "--deployer-attribute-file",
+    default=None,
+    show_default=True,
+    type=str,
+    help="Write the metadata and pathspec of this run to the file specified.\nUsed internally for Metaflow's Deployer API.",
+    hidden=True,
+)
 @click.pass_obj
-def trigger(obj, run_id_file=None, **kwargs):
+def trigger(obj, run_id_file=None, deployer_attribute_file=None, **kwargs):
     def _convert_value(param):
         # Swap `-` with `_` in parameter name to match click's behavior
         val = kwargs.get(param.name.replace("-", "_").lower())
@@ -466,6 +549,17 @@ def trigger(obj, run_id_file=None, **kwargs):
         with open(run_id_file, "w") as f:
             f.write(str(run_id))
 
+    if deployer_attribute_file:
+        with open(deployer_attribute_file, "w") as f:
+            json.dump(
+                {
+                    "name": obj.state_machine_name,
+                    "metadata": obj.metadata.metadata_str(),
+                    "pathspec": "/".join((obj.flow.name, run_id)),
+                },
+                f,
+            )
+
     obj.echo(
         "Workflow *{name}* triggered on AWS Step Functions "
         "(run-id *{run_id}*).".format(name=obj.state_machine_name, run_id=run_id),

metaflow/plugins/aws/step_functions/step_functions_decorator.py

@@ -1,9 +1,8 @@
-import json
 import os
 import time
 
 from metaflow.decorators import StepDecorator
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 
 from .dynamo_db_client import DynamoDbClient
 

metaflow/plugins/aws/step_functions/step_functions_deployer.py

@@ -0,0 +1,97 @@
+from typing import Any, ClassVar, Dict, Optional, TYPE_CHECKING, Type
+
+from metaflow.runner.deployer_impl import DeployerImpl
+
+if TYPE_CHECKING:
+    import metaflow.plugins.aws.step_functions.step_functions_deployer_objects
+
+
+class StepFunctionsDeployer(DeployerImpl):
+    """
+    Deployer implementation for AWS Step Functions.
+
+    Parameters
+    ----------
+    name : str, optional, default None
+        State Machine name. The flow name is used instead if this option is not specified.
+    """
+
+    TYPE: ClassVar[Optional[str]] = "step-functions"
+
+    def __init__(self, deployer_kwargs: Dict[str, str], **kwargs):
+        """
+        Initialize the StepFunctionsDeployer.
+
+        Parameters
+        ----------
+        deployer_kwargs : Dict[str, str]
+            The deployer-specific keyword arguments.
+        **kwargs : Any
+            Additional arguments to pass to the superclass constructor.
+        """
+        self._deployer_kwargs = deployer_kwargs
+        super().__init__(**kwargs)
+
+    @property
+    def deployer_kwargs(self) -> Dict[str, Any]:
+        return self._deployer_kwargs
+
+    @staticmethod
+    def deployed_flow_type() -> (
+        Type[
+            "metaflow.plugins.aws.step_functions.step_functions_deployer_objects.StepFunctionsDeployedFlow"
+        ]
+    ):
+        from .step_functions_deployer_objects import StepFunctionsDeployedFlow
+
+        return StepFunctionsDeployedFlow
+
+    def create(
+        self, **kwargs
+    ) -> "metaflow.plugins.aws.step_functions.step_functions_deployer_objects.StepFunctionsDeployedFlow":
+        """
+        Create a new AWS Step Functions State Machine deployment.
+
+        Parameters
+        ----------
+        authorize : str, optional, default None
+            Authorize using this production token. Required when re-deploying an existing flow
+            for the first time. The token is cached in METAFLOW_HOME.
+        generate_new_token : bool, optional, default False
+            Generate a new production token for this flow. Moves the production flow to a new namespace.
+        given_token : str, optional, default None
+            Use the given production token for this flow. Moves the production flow to the given namespace.
+        tags : List[str], optional, default None
+            Annotate all objects produced by AWS Step Functions runs with these tags.
+        user_namespace : str, optional, default None
+            Change the namespace from the default (production token) to the given tag.
+        only_json : bool, optional, default False
+            Only print out JSON sent to AWS Step Functions without deploying anything.
+        max_workers : int, optional, default 100
+            Maximum number of parallel processes.
+        workflow_timeout : int, optional, default None
+            Workflow timeout in seconds.
+        log_execution_history : bool, optional, default False
+            Log AWS Step Functions execution history to AWS CloudWatch Logs log group.
+        use_distributed_map : bool, optional, default False
+            Use AWS Step Functions Distributed Map instead of Inline Map for defining foreach
+            tasks in Amazon State Language.
+        compress_state_machine : bool, optional, default False
+            Compress AWS Step Functions state machine to fit within the 8K limit.
+
+        deployer_attribute_file : str, optional, default None
+            Write the workflow name to the specified file. Used internally for Metaflow's Deployer API.
+
+        Returns
+        -------
+        StepFunctionsDeployedFlow
+            The Flow deployed to AWS Step Functions.
+        """
+        from .step_functions_deployer_objects import StepFunctionsDeployedFlow
+
+        return self._create(StepFunctionsDeployedFlow, **kwargs)
+
+
+_addl_stubgen_modules = [
+    "metaflow.plugins.aws.step_functions.step_functions_deployer_objects"
+]