From 345e0f4171c8472e57e3649e27ca8d0a1cdfb41f Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Mon, 9 Jan 2023 08:27:33 +0000 Subject: [PATCH 1/4] Support for multi-flow decorators. --- metaflow/decorators.py | 34 ++++++++++++------- metaflow/plugins/airflow/airflow.py | 1 + metaflow/plugins/argo/argo_workflows.py | 1 + .../aws/step_functions/step_functions.py | 1 + .../plugins/conda/conda_step_decorator.py | 2 +- 5 files changed, 26 insertions(+), 13 deletions(-) diff --git a/metaflow/decorators.py b/metaflow/decorators.py index 8723367b79b..b13b927825e 100644 --- a/metaflow/decorators.py +++ b/metaflow/decorators.py @@ -108,6 +108,8 @@ class Decorator(object): name = "NONAME" defaults = {} + # `allow_multiple` allows setting many decorators of the same type to a step/flow. + allow_multiple = False def __init__(self, attributes=None, statically_defined=False): self.attributes = self.defaults.copy() @@ -255,9 +257,6 @@ class MyDecorator(StepDecorator): pass them around with every lifecycle call. """ - # `allow_multiple` allows setting many decorators of the same type to a step. - allow_multiple = False - def step_init( self, flow, graph, step_name, decorators, environment, flow_datastore, logger ): @@ -402,13 +401,15 @@ def _base_flow_decorator(decofunc, *args, **kwargs): cls = args[0] if isinstance(cls, type) and issubclass(cls, FlowSpec): # flow decorators add attributes in the class dictionary, - # _flow_decorators. - if decofunc.name in cls._flow_decorators: + # _flow_decorators. _flow_decorators is of type `{key:[decos]}` + if decofunc.name in cls._flow_decorators and not decofunc.allow_multiple: raise DuplicateFlowDecoratorException(decofunc.name) else: - cls._flow_decorators[decofunc.name] = decofunc( - attributes=kwargs, statically_defined=True - ) + deco_instance = decofunc(attributes=kwargs, statically_defined=True) + if decofunc.name not in cls._flow_decorators: + cls._flow_decorators[decofunc.name] = [deco_instance] + else: + cls._flow_decorators[decofunc.name].append(deco_instance) else: raise BadFlowDecoratorException(decofunc.name) return cls @@ -503,11 +504,20 @@ def _attach_decorators_to_step(step, decospecs): def _init_flow_decorators( flow, graph, environment, flow_datastore, metadata, logger, echo, deco_options ): + # Since all flow decorators are stored as `{key:[deco]}` we iterate through each of them. for deco in flow._flow_decorators.values(): - opts = {option: deco_options[option] for option in deco.options} - deco.flow_init( - flow, graph, environment, flow_datastore, metadata, logger, echo, opts - ) + for rd in deco: + opts = {option: deco_options[option] for option in rd.options} + rd.flow_init( + flow, + graph, + environment, + flow_datastore, + metadata, + logger, + echo, + opts, + ) def _init_step_decorators(flow, graph, environment, flow_datastore, logger): diff --git a/metaflow/plugins/airflow/airflow.py b/metaflow/plugins/airflow/airflow.py index 5480a79c59e..02a71d856ab 100644 --- a/metaflow/plugins/airflow/airflow.py +++ b/metaflow/plugins/airflow/airflow.py @@ -140,6 +140,7 @@ def _get_schedule(self): schedule = self.flow._flow_decorators.get("schedule") if not schedule: return None + schedule = schedule[0] if schedule.attributes["cron"]: return schedule.attributes["cron"] elif schedule.attributes["weekly"]: diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index 1730b493248..caa61403e70 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -174,6 +174,7 @@ def trigger(cls, name, parameters=None): def _cron(self): schedule = self.flow._flow_decorators.get("schedule") + schedule = schedule[0] if schedule: # Remove the field "Year" if it exists return " ".join(schedule.schedule.split()[:5]), schedule.timezone diff --git a/metaflow/plugins/aws/step_functions/step_functions.py b/metaflow/plugins/aws/step_functions/step_functions.py index 8c9584ed706..370d4adec93 100644 --- a/metaflow/plugins/aws/step_functions/step_functions.py +++ b/metaflow/plugins/aws/step_functions/step_functions.py @@ -328,6 +328,7 @@ def _visit(node, workflow, exit_node=None): def _cron(self): schedule = self.flow._flow_decorators.get("schedule") + schedule = schedule[0] if schedule: if schedule.timezone is not None: raise StepFunctionsException( diff --git a/metaflow/plugins/conda/conda_step_decorator.py b/metaflow/plugins/conda/conda_step_decorator.py index 6896e293c25..818e312a115 100644 --- a/metaflow/plugins/conda/conda_step_decorator.py +++ b/metaflow/plugins/conda/conda_step_decorator.py @@ -63,7 +63,7 @@ class CondaStepDecorator(StepDecorator): def _get_base_attributes(self): if "conda_base" in self.flow._flow_decorators: - return self.flow._flow_decorators["conda_base"].attributes + return self.flow._flow_decorators["conda_base"][0].attributes return self.defaults def _python_version(self): From e46c51b0c63b4a83dcc71f70f688085df3802b6e Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Fri, 27 Jan 2023 19:47:01 +0000 Subject: [PATCH 2/4] Refactor option setting from flow decorators. --- metaflow/decorators.py | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/metaflow/decorators.py b/metaflow/decorators.py index b13b927825e..b8bdcaaedcb 100644 --- a/metaflow/decorators.py +++ b/metaflow/decorators.py @@ -505,10 +505,28 @@ def _init_flow_decorators( flow, graph, environment, flow_datastore, metadata, logger, echo, deco_options ): # Since all flow decorators are stored as `{key:[deco]}` we iterate through each of them. - for deco in flow._flow_decorators.values(): - for rd in deco: - opts = {option: deco_options[option] for option in rd.options} - rd.flow_init( + for decorators in flow._flow_decorators.values(): + # First resolve the `options` for the flow decorator. + # Options are passed from cli. + # For example `@project` can take a `--name` / `--branch` from the cli as options. + deco_flow_init_options = {} + deco = decorators[0] + # If a flow decorator allow multiple of same type then we don't allow multiple options for it. + if deco.allow_multiple: + if len(deco.options) > 0: + raise MetaflowException( + "Flow decorator `@%s` has multiple options, which is not allowed. " + "Please ensure the FlowDecorator `%s` has no options since flow decorators with " + "`allow_mutiple=True` are not allowed to have options" + % (deco.name, deco.__class__.__name__) + ) + else: + # Each "non-multiple" flow decorator is only allowed to have one set of options + deco_flow_init_options = { + option: deco_options[option] for option in deco.options + } + for deco in decorators: + deco.flow_init( flow, graph, environment, @@ -516,7 +534,7 @@ def _init_flow_decorators( metadata, logger, echo, - opts, + deco_flow_init_options, ) From 48197c056004df241dbc83619aba9e0d0d47a21c Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Fri, 27 Jan 2023 19:52:38 +0000 Subject: [PATCH 3/4] refactor logic for simplicity --- metaflow/decorators.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/metaflow/decorators.py b/metaflow/decorators.py index b8bdcaaedcb..b3f88df56d5 100644 --- a/metaflow/decorators.py +++ b/metaflow/decorators.py @@ -406,10 +406,7 @@ def _base_flow_decorator(decofunc, *args, **kwargs): raise DuplicateFlowDecoratorException(decofunc.name) else: deco_instance = decofunc(attributes=kwargs, statically_defined=True) - if decofunc.name not in cls._flow_decorators: - cls._flow_decorators[decofunc.name] = [deco_instance] - else: - cls._flow_decorators[decofunc.name].append(deco_instance) + cls._flow_decorators.setdefault(decofunc.name, []).append(deco_instance) else: raise BadFlowDecoratorException(decofunc.name) return cls From 75c690b11dc943f9abbbcf6dde6debed32364747 Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Fri, 27 Jan 2023 21:25:37 +0000 Subject: [PATCH 4/4] Fixed bug due to fat-fingered indent --- metaflow/decorators.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/metaflow/decorators.py b/metaflow/decorators.py index b3f88df56d5..ebdd77ebd7e 100644 --- a/metaflow/decorators.py +++ b/metaflow/decorators.py @@ -517,11 +517,11 @@ def _init_flow_decorators( "`allow_mutiple=True` are not allowed to have options" % (deco.name, deco.__class__.__name__) ) - else: - # Each "non-multiple" flow decorator is only allowed to have one set of options - deco_flow_init_options = { - option: deco_options[option] for option in deco.options - } + else: + # Each "non-multiple" flow decorator is only allowed to have one set of options + deco_flow_init_options = { + option: deco_options[option] for option in deco.options + } for deco in decorators: deco.flow_init( flow,