diff --git a/metaflow/cli.py b/metaflow/cli.py index 1999c3c29f..4923902365 100644 --- a/metaflow/cli.py +++ b/metaflow/cli.py @@ -461,7 +461,7 @@ def start( ) if all_decospecs: decorators._attach_decorators(ctx.obj.flow, all_decospecs) - decorators._init(ctx.obj.flow, only_non_static=True) + decorators._init(ctx.obj.flow) # Regenerate graph if we attached more decorators ctx.obj.graph = FlowGraph(ctx.obj.flow.__class__) diff --git a/metaflow/cli_components/run_cmds.py b/metaflow/cli_components/run_cmds.py index ec31e2d3aa..4d01349c97 100644 --- a/metaflow/cli_components/run_cmds.py +++ b/metaflow/cli_components/run_cmds.py @@ -40,7 +40,7 @@ def before_run(obj, tags, decospecs): ) if all_decospecs: decorators._attach_decorators(obj.flow, all_decospecs) - decorators._init(obj.flow, only_non_static=True) + decorators._init(obj.flow) obj.graph = FlowGraph(obj.flow.__class__) obj.check(obj.graph, obj.flow, obj.environment, pylint=obj.pylint) diff --git a/metaflow/cli_components/step_cmd.py b/metaflow/cli_components/step_cmd.py index 9302870fec..bb4df976e7 100644 --- a/metaflow/cli_components/step_cmd.py +++ b/metaflow/cli_components/step_cmd.py @@ -138,7 +138,7 @@ def step( if decospecs: decorators._attach_decorators_to_step(func, decospecs) - decorators._init(ctx.obj.flow, only_non_static=True) + decorators._init(ctx.obj.flow) step_kwargs = ctx.params # Remove argument `step_name` from `step_kwargs`. diff --git a/metaflow/decorators.py b/metaflow/decorators.py index 35ce578db5..23c3222a1d 100644 --- a/metaflow/decorators.py +++ b/metaflow/decorators.py @@ -27,6 +27,11 @@ unicode = str basestring = str +# Contains the decorators on which _init was called. We want to ensure it is called +# only once on each decorator and, as the _init() function below can be called in +# several places, we need to track which decorator had their init function called +_inited_decorators = set() + class BadStepDecoratorException(MetaflowException): headline = "Syntax error" @@ -553,12 +558,16 @@ def _attach_decorators_to_step(step, decospecs): def _init(flow, only_non_static=False): for decorators in flow._flow_decorators.values(): for deco in decorators: - if not only_non_static or not deco.statically_defined: - deco.init() + if deco in _inited_decorators: + continue + deco.init() + _inited_decorators.add(deco) for flowstep in flow: for deco in flowstep.decorators: - if not only_non_static or not deco.statically_defined: - deco.init() + if deco in _inited_decorators: + continue + deco.init() + _inited_decorators.add(deco) def _init_flow_decorators( diff --git a/metaflow/plugins/airflow/airflow_cli.py b/metaflow/plugins/airflow/airflow_cli.py index 33e0e0f48d..63b96b46a2 100644 --- a/metaflow/plugins/airflow/airflow_cli.py +++ b/metaflow/plugins/airflow/airflow_cli.py @@ -283,7 +283,7 @@ def make_flow( ): # Attach @kubernetes. decorators._attach_decorators(obj.flow, [KubernetesDecorator.name]) - decorators._init(obj.flow, only_non_static=True) + decorators._init(obj.flow) decorators._init_step_decorators( obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger diff --git a/metaflow/plugins/argo/argo_workflows_cli.py b/metaflow/plugins/argo/argo_workflows_cli.py index 734c50d21e..b96ffbbd69 100644 --- a/metaflow/plugins/argo/argo_workflows_cli.py +++ b/metaflow/plugins/argo/argo_workflows_cli.py @@ -470,7 +470,7 @@ def make_flow( decorators._attach_decorators( obj.flow, [KubernetesDecorator.name, EnvironmentDecorator.name] ) - decorators._init(obj.flow, only_non_static=True) + decorators._init(obj.flow) decorators._init_step_decorators( obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger diff --git a/metaflow/plugins/aws/step_functions/step_functions_cli.py b/metaflow/plugins/aws/step_functions/step_functions_cli.py index 373b5daef1..63682b5e7d 100644 --- a/metaflow/plugins/aws/step_functions/step_functions_cli.py +++ b/metaflow/plugins/aws/step_functions/step_functions_cli.py @@ -326,7 +326,7 @@ def make_flow( # Attach AWS Batch decorator to the flow decorators._attach_decorators(obj.flow, [BatchDecorator.name]) - decorators._init(obj.flow, only_non_static=True) + decorators._init(obj.flow) decorators._init_step_decorators( obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger ) diff --git a/metaflow/plugins/pypi/pypi_decorator.py b/metaflow/plugins/pypi/pypi_decorator.py index 3f650281dc..1b16967ced 100644 --- a/metaflow/plugins/pypi/pypi_decorator.py +++ b/metaflow/plugins/pypi/pypi_decorator.py @@ -140,7 +140,7 @@ def flow_init( from metaflow import decorators decorators._attach_decorators(flow, ["pypi"]) - decorators._init(flow, only_non_static=True) + decorators._init(flow) # @pypi uses a conda environment to create a virtual environment. # The conda environment can be created through micromamba.