Fix case where the decorator's init function could be called twice. (N…

…etflix#2167)

In some cases, specifically when using argo/airflow/sfn, the decorator's
init function could be called twice which caused issues with the conda/pypi
decorator.

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__)
 

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)

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`.

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(

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

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

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
     )

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.