Re-add backward compatible _flow_decorators

Author: romain-intel

metaflow/flowspec.py

@@ -326,6 +326,11 @@ def script_name(self) -> str:
             fname = fname[:-1]
         return os.path.basename(fname)
 
+    @property
+    def _flow_decorators(self):
+        # Backward compatible method to access flow decorators
+        return self._flow_state[FlowStateItems.FLOW_DECORATORS]
+
     @classmethod
     def _check_parameters(cls, config_parameters=False):
         seen = set()

metaflow/plugins/airflow/airflow.py

@@ -10,7 +10,6 @@
 from metaflow import current
 from metaflow.decorators import flow_decorators
 from metaflow.exception import MetaflowException
-from metaflow.flowspec import FlowStateItems
 from metaflow.includefile import FilePathClass
 from metaflow.metaflow_config import (
     AIRFLOW_KUBERNETES_CONN_ID,
@@ -151,7 +150,7 @@ def save_deployment_token(cls, owner, name, token, flow_datastore):
     def _get_schedule(self):
         # Using the cron presets provided here :
         # https://airflow.apache.org/docs/apache-airflow/stable/dag-run.html?highlight=schedule%20interval#cron-presets
-        schedule = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS].get("schedule")
+        schedule = self.flow._flow_decorators.get("schedule")
         if not schedule:
             return None
         schedule = schedule[0]
@@ -634,11 +633,10 @@ def _step_cli(self, node, paths, code_package_url, user_code_retries):
         return cmds
 
     def _collect_flow_sensors(self):
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
         decos_lists = [
-            flow_decos.get(s.name)
+            self.flow._flow_decorators.get(s.name)
             for s in SUPPORTED_SENSORS
-            if flow_decos.get(s.name) is not None
+            if self.flow._flow_decorators.get(s.name) is not None
         ]
         af_tasks = [deco.create_task() for decos in decos_lists for deco in decos]
         if len(af_tasks) > 0:
@@ -652,14 +650,15 @@ def _contains_foreach(self):
         return False
 
     def compile(self):
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
-        if flow_decos.get("trigger") or flow_decos.get("trigger_on_finish"):
+        if self.flow._flow_decorators.get("trigger") or self.flow._flow_decorators.get(
+            "trigger_on_finish"
+        ):
             raise AirflowException(
                 "Deploying flows with @trigger or @trigger_on_finish decorator(s) "
                 "to Airflow is not supported currently."
             )
 
-        if flow_decos.get("exit_hook"):
+        if self.flow._flow_decorators.get("exit_hook"):
             raise AirflowException(
                 "Deploying flows with the @exit_hook decorator "
                 "to Airflow is not currently supported."

metaflow/plugins/airflow/airflow_cli.py

@@ -7,7 +7,6 @@
 from metaflow import current, decorators
 from metaflow._vendor import click
 from metaflow.exception import MetaflowException, MetaflowInternalError
-from metaflow.flowspec import FlowStateItems
 from metaflow.metaflow_config import FEAT_ALWAYS_UPLOAD_CODE_PACKAGE
 from metaflow.package import MetaflowPackage
 from metaflow.plugins.aws.step_functions.production_token import (
@@ -418,7 +417,7 @@ def _validate_workflow(flow, graph, flow_datastore, metadata, workflow_timeout):
             )
         )
 
-    schedule = flow._flow_state[FlowStateItems.FLOW_DECORATORS].get("schedule")
+    schedule = flow._flow_decorators.get("schedule")
     if not schedule:
         return
 

metaflow/plugins/argo/argo_workflows.py

@@ -12,7 +12,6 @@
 from metaflow import JSONType, current
 from metaflow.decorators import flow_decorators
 from metaflow.exception import MetaflowException
-from metaflow.flowspec import FlowStateItems
 from metaflow.graph import FlowGraph
 from metaflow.includefile import FilePathClass
 from metaflow.metaflow_config import (
@@ -434,7 +433,7 @@ def _base_kubernetes_annotations(self):
         return annotations
 
     def _get_schedule(self):
-        schedule = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS].get("schedule")
+        schedule = self.flow._flow_decorators.get("schedule")
         if schedule:
             # Remove the field "Year" if it exists
             schedule = schedule[0]
@@ -460,15 +459,14 @@ def schedule(self):
 
     def trigger_explanation(self):
         # Trigger explanation for cron workflows
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
-        if flow_decos.get("schedule"):
+        if self.flow._flow_decorators.get("schedule"):
             return (
                 "This workflow triggers automatically via the CronWorkflow *%s*."
                 % self.name
             )
 
         # Trigger explanation for @trigger
-        elif flow_decos.get("trigger"):
+        elif self.flow._flow_decorators.get("trigger"):
             return (
                 "This workflow triggers automatically when the upstream %s "
                 "is/are published."
@@ -478,7 +476,7 @@ def trigger_explanation(self):
             )
 
         # Trigger explanation for @trigger_on_finish
-        elif flow_decos.get("trigger_on_finish"):
+        elif self.flow._flow_decorators.get("trigger_on_finish"):
             return (
                 "This workflow triggers automatically when the upstream %s succeed(s)"
                 % self.list_to_prose(
@@ -541,10 +539,7 @@ def get_execution(cls, name):
 
     def _process_parameters(self):
         parameters = {}
-        has_schedule = (
-            self.flow._flow_state[FlowStateItems.FLOW_DECORATORS].get("schedule")
-            is not None
-        )
+        has_schedule = self.flow._flow_decorators.get("schedule") is not None
         seen = set()
         for var, param in self.flow._get_parameters():
             # Throw an exception if the parameter is specified twice.
@@ -630,9 +625,10 @@ def _process_triggers(self):
         # Impute triggers for Argo Workflow Template specified through @trigger and
         # @trigger_on_finish decorators
 
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
         # Disallow usage of @trigger and @trigger_on_finish together for now.
-        if flow_decos.get("trigger") and flow_decos.get("trigger_on_finish"):
+        if self.flow._flow_decorators.get("trigger") and self.flow._flow_decorators.get(
+            "trigger_on_finish"
+        ):
             raise ArgoWorkflowsException(
                 "Argo Workflows doesn't support both *@trigger* and "
                 "*@trigger_on_finish* decorators concurrently yet. Use one or the "
@@ -642,7 +638,7 @@ def _process_triggers(self):
         options = None
 
         # @trigger decorator
-        if flow_decos.get("trigger"):
+        if self.flow._flow_decorators.get("trigger"):
             # Parameters are not duplicated, and exist in the flow. Additionally,
             # convert them to lower case since Metaflow parameters are case
             # insensitive.
@@ -655,7 +651,7 @@ def _process_triggers(self):
                     if not param.IS_CONFIG_PARAMETER
                 ]
             )
-            trigger_deco = flow_decos.get("trigger")[0]
+            trigger_deco = self.flow._flow_decorators.get("trigger")[0]
             trigger_deco.format_deploytime_value()
             for event in trigger_deco.triggers:
                 parameters = {}
@@ -687,17 +683,19 @@ def _process_triggers(self):
                     parameters[key.lower()] = value
                 event["parameters"] = parameters
                 event["type"] = "event"
-            triggers.extend(flow_decos.get("trigger")[0].triggers)
+            triggers.extend(self.flow._flow_decorators.get("trigger")[0].triggers)
 
             # Set automatic parameter mapping iff only a single event dependency is
             # specified with no explicit parameter mapping.
             if len(triggers) == 1 and not triggers[0].get("parameters"):
                 triggers[0]["parameters"] = dict(zip(params, params))
-            options = flow_decos.get("trigger")[0].options
+            options = self.flow._flow_decorators.get("trigger")[0].options
 
         # @trigger_on_finish decorator
-        if flow_decos.get("trigger_on_finish"):
-            trigger_on_finish_deco = flow_decos.get("trigger_on_finish")[0]
+        if self.flow._flow_decorators.get("trigger_on_finish"):
+            trigger_on_finish_deco = self.flow._flow_decorators.get(
+                "trigger_on_finish"
+            )[0]
             trigger_on_finish_deco.format_deploytime_value()
             for event in trigger_on_finish_deco.triggers:
                 # Actual filters are deduced here since we don't have access to
@@ -736,7 +734,7 @@ def _process_triggers(self):
                         "flow": event["flow"],
                     }
                 )
-            options = flow_decos.get("trigger_on_finish")[0].options
+            options = self.flow._flow_decorators.get("trigger_on_finish")[0].options
 
         for event in triggers:
             # Assign a sanitized name since we need this at many places to please
@@ -2805,9 +2803,7 @@ def _lifecycle_hooks(self):
             hooks.append(self._pager_duty_change_template())
             hooks.append(self._incident_io_change_template())
 
-        exit_hook_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS].get(
-            "exit_hook", []
-        )
+        exit_hook_decos = self.flow._flow_decorators.get("exit_hook", [])
 
         for deco in exit_hook_decos:
             hooks.extend(self._lifecycle_hook_from_deco(deco))

metaflow/plugins/aws/step_functions/step_functions.py

@@ -9,7 +9,6 @@
 from metaflow import R
 from metaflow.decorators import flow_decorators
 from metaflow.exception import MetaflowException
-from metaflow.flowspec import FlowStateItems
 from metaflow.metaflow_config import (
     EVENTS_SFN_ACCESS_IAM_ROLE,
     S3_ENDPOINT_URL,
@@ -303,14 +302,15 @@ def get_execution(cls, state_machine_name, name):
             raise StepFunctionsException(repr(e))
 
     def _compile(self):
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
-        if flow_decos.get("trigger") or flow_decos.get("trigger_on_finish"):
+        if self.flow._flow_decorators.get("trigger") or self.flow._flow_decorators.get(
+            "trigger_on_finish"
+        ):
             raise StepFunctionsException(
                 "Deploying flows with @trigger or @trigger_on_finish decorator(s) "
                 "to AWS Step Functions is not supported currently."
             )
 
-        if flow_decos.get("exit_hook"):
+        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."
@@ -484,8 +484,7 @@ def _visit(node, workflow, exit_node=None):
         return _visit(self.graph["start"], workflow)
 
     def _cron(self):
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
-        schedule = flow_decos.get("schedule")
+        schedule = self.flow._flow_decorators.get("schedule")
         if schedule:
             schedule = schedule[0]
             if schedule.timezone is not None:

metaflow/plugins/pypi/conda_decorator.py

@@ -5,7 +5,6 @@
 import tempfile
 
 from metaflow.decorators import FlowDecorator, StepDecorator
-from metaflow.flowspec import FlowStateItems
 from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
 from metaflow.packaging_sys import ContentType
@@ -83,9 +82,8 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
         self.datastore = flow_datastore
 
         # Support flow-level decorator.
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
-        if "conda_base" in flow_decos:
-            conda_base = flow_decos["conda_base"][0]
+        if "conda_base" in self.flow._flow_decorators:
+            conda_base = self.flow._flow_decorators["conda_base"][0]
             super_attributes = conda_base.attributes
             self.attributes["packages"] = {
                 **super_attributes["packages"],

metaflow/plugins/pypi/pypi_decorator.py

@@ -1,5 +1,4 @@
 from metaflow.decorators import FlowDecorator, StepDecorator
-from metaflow.flowspec import FlowStateItems
 from metaflow.metaflow_environment import InvalidEnvironmentException
 
 
@@ -41,9 +40,8 @@ def step_init(self, flow, graph, step, decos, environment, flow_datastore, logge
         self.step = step
 
         # Support flow-level decorator
-        flow_decos = self.flow._flow_state[FlowStateItems.FLOW_DECORATORS]
-        if "pypi_base" in flow_decos:
-            pypi_base = flow_decos["pypi_base"][0]
+        if "pypi_base" in self.flow._flow_decorators:
+            pypi_base = self.flow._flow_decorators["pypi_base"][0]
             super_attributes = pypi_base.attributes
             self._attributes_with_user_values.update(
                 pypi_base._attributes_with_user_values