Addressed an issue with class members (#804)

romain-intel · tuulos · web-flow · commit ac7d7df6ce4b · 2021-11-03T18:28:27.000-07:00
* Addressed an issue with class members

Class members that were not parameters caused an error (regression
in 2.4.0). This addresses this issue.

This patch also clarifies and makes explicit that modifying class
variables is not safe; an error message is now raised.

* Avoid persisting class variables multiple times.

This treats class variables more like parameters (constants) and
persists them once in the renamed persist_constants (used to be
persist_parameters).

* add constants test

* Make test pass (hopefully)

* add CLASS_VARS to tests

Co-authored-by: Ville Tuulos &lt;tuulos@gmail.com&gt;
diff --git a/docs/lifecycle.dot b/docs/lifecycle.dot
@@ -53,7 +53,7 @@ digraph Metaflow {
         command_run         [label="{command|run}", fillcolor=tan]
         new_run_id          [label="{metadata|new_run_id}", fillcolor=lightgoldenrod1]
         runtime_init        [label="{decorator|runtime_init}", fillcolor=lightblue2]
-        local_params        [label="{runtime|persist_parameters}", fillcolor=lightpink2]
+        local_params        [label="{runtime|persist_constants}", fillcolor=lightpink2]
         start_run_heartbeat [label="{metadata|start_run_heartbeat}", fillcolor=lightgoldenrod1]
         schedule_local_task [shape="circle", label="Schedule\nTask", width=1, fillcolor=grey78]
         runtime_finished    [label="{decorator|runtime_finished}", fillcolor=lightblue2]
@@ -99,7 +99,7 @@ digraph Metaflow {
         stepfunctions_run             [label="{AWS Step Functions|start_execution}", fillcolor=lightpink2]
         stepfunctions_bootstrap_batch [shape="circle", label="Bootstrap\nAWS Batch", width=1, fillcolor=grey78]
         stepfunctions_init            [label="{command|init}" fillcolor=tan]
-        stepfunctions_params          [label="{runtime|persist_parameters}", fillcolor=lightpink2]
+        stepfunctions_params          [label="{runtime|persist_constants}", fillcolor=lightpink2]
         stepfunctions_task            [shape="circle", label="Execute\nTask", width=1, fillcolor=grey78]
     }
 
diff --git a/metaflow/cli.py b/metaflow/cli.py
@@ -592,8 +592,8 @@ def init(obj, run_id=None, task_id=None, tags=None, **kwargs):
         obj.monitor,
         run_id=run_id,
     )
-    parameters.set_parameters(obj.flow, kwargs)
-    runtime.persist_parameters(task_id=task_id)
+    obj.flow._set_constants(kwargs)
+    runtime.persist_constants(task_id=task_id)
 
 
 def common_run_options(func):
@@ -711,7 +711,7 @@ def resume(
         max_num_splits=max_num_splits,
         max_log_size=max_log_size * 1024 * 1024,
     )
-    runtime.persist_parameters()
+    runtime.persist_constants()
     runtime.execute()
 
     write_run_id(run_id_file, runtime.run_id)
@@ -766,8 +766,8 @@ def run(
     write_latest_run_id(obj, runtime.run_id)
     write_run_id(run_id_file, runtime.run_id)
 
-    parameters.set_parameters(obj.flow, kwargs)
-    runtime.persist_parameters()
+    obj.flow._set_constants(kwargs)
+    runtime.persist_constants()
     runtime.execute()
 
 
diff --git a/metaflow/datastore/task_datastore.py b/metaflow/datastore/task_datastore.py
@@ -643,7 +643,7 @@ def persist(self, flow):
         for var in dir(flow):
             if var.startswith("__") or var in flow._EPHEMERAL:
                 continue
-            # Skip over properties of the class (Parameters)
+            # Skip over properties of the class (Parameters or class variables)
             if hasattr(flow.__class__, var) and isinstance(
                 getattr(flow.__class__, var), property
             ):
diff --git a/metaflow/flowspec.py b/metaflow/flowspec.py
@@ -3,6 +3,7 @@
 import sys
 import inspect
 import traceback
+from types import FunctionType, MethodType
 
 from . import cmd_with_io
 from .parameters import Parameter
@@ -106,6 +107,44 @@ def script_name(self):
             fname = fname[:-1]
         return os.path.basename(fname)
 
+    def _set_constants(self, kwargs):
+        # Persist values for parameters and other constants (class level variables)
+        # only once. This method is called before persist_constants is called to
+        # persist all values set using setattr
+        seen = set()
+        for var, param in self._get_parameters():
+            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)
+        seen.clear()
+        self._success = True
+
+        for var, param in self._get_parameters():
+            seen.add(var)
+            val = kwargs[param.name.replace("-", "_").lower()]
+            # Support for delayed evaluation of parameters. This is used for
+            # includefile in particular
+            if callable(val):
+                val = val()
+            val = val.split(param.separator) if val and param.separator else val
+            setattr(self, var, val)
+
+        # Do the same for class variables which will be forced constant as modifications
+        # to them don't propagate well since we create a new process for each step and
+        # re-read the flow file
+        for var in dir(self.__class__):
+            if var[0] == "_" or var in self._NON_PARAMETERS or var in seen:
+                continue
+            val = getattr(self.__class__, var)
+            if isinstance(val, (MethodType, FunctionType, property, type)):
+                continue
+            setattr(self, var, val)
+
     def _get_parameters(self):
         for var in dir(self):
             if var[0] == "_" or var in self._NON_PARAMETERS:
diff --git a/metaflow/includefile.py b/metaflow/includefile.py
@@ -232,7 +232,7 @@ class FilePathClass(click.ParamType):
     #    + If the value is already such a string, nothing happens and it returns that same value
     #    + If the value is a LocalFile, it will persist the local file and return the path
     #      of the persisted file
-    #  - The artifact will be persisted prior to any run (for non-scheduled runs through persist_parameters)
+    #  - The artifact will be persisted prior to any run (for non-scheduled runs through persist_constants)
     #    + This will therefore persist a simple string
     #  - When the parameter is loaded again, the load_parameter in the IncludeFile class will get called
     #    which will download and return the bytes of the persisted file.
diff --git a/metaflow/parameters.py b/metaflow/parameters.py
@@ -250,26 +250,3 @@ def wrapper(cmd):
         return cmd
 
     return wrapper
-
-
-def set_parameters(flow, kwargs):
-    seen = set()
-    for var, param in flow._get_parameters():
-        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)
-
-    flow._success = True
-    for var, param in flow._get_parameters():
-        val = kwargs[param.name.replace("-", "_").lower()]
-        # Support for delayed evaluation of parameters. This is used for
-        # includefile in particular
-        if callable(val):
-            val = val()
-        val = val.split(param.separator) if val and param.separator else val
-        setattr(flow, var, val)
diff --git a/metaflow/runtime.py b/metaflow/runtime.py
@@ -176,7 +176,7 @@ def _new_task(self, step, input_paths=None, **kwargs):
     def run_id(self):
         return self._run_id
 
-    def persist_parameters(self, task_id=None):
+    def persist_constants(self, task_id=None):
         task = self._new_task("_parameters", task_id=task_id)
         if not task.is_cloned:
             task.persist(self._flow)
@@ -614,7 +614,7 @@ def __init__(
         if task_id is None:
             task_id = str(metadata.new_task_id(run_id, step))
         else:
-            # task_id is preset only by persist_parameters() or control tasks.
+            # task_id is preset only by persist_constants() or control tasks.
             if ubf_context == UBF_CONTROL:
                 metadata.register_task_id(
                     run_id, step, task_id, 0, sys_tags=[CONTROL_TASK_TAG]
diff --git a/metaflow/task.py b/metaflow/task.py
@@ -3,6 +3,8 @@
 import os
 import time
 
+from types import MethodType, FunctionType
+
 from .metaflow_config import MAX_ATTEMPTS
 from .metadata import MetaDatum
 from .datastore import Inputs, TaskDataStoreSet
@@ -53,14 +55,18 @@ def _exec_step_function(self, step_function, input_obj=None):
             step_function(input_obj)
 
     def _init_parameters(self, parameter_ds, passdown=True):
+        def set_cls_var(_, __):
+            raise AttributeError("Flow level attributes are not modifiable")
+
+        cls = self.flow.__class__
         # overwrite Parameters in the flow object
         vars = []
         for var, param in self.flow._get_parameters():
             # make the parameter a read-only property
             # note x=x binds the current value of x to the closure
             def property_setter(
                 _,
-                cls=self.flow.__class__,
+                cls=cls,
                 param=param,
                 var=var,
                 parameter_ds=parameter_ds,
@@ -69,11 +75,26 @@ def property_setter(
                 setattr(cls, var, property(fget=lambda _, val=v: val))
                 return v
 
-            setattr(self.flow.__class__, var, property(fget=property_setter))
+            setattr(cls, var, property(fget=property_setter))
             vars.append(var)
+
+        param_only_vars = list(vars)
+        # make class-level values read-only to be more consistent across steps in a flow
+        # they are also only persisted once and so we similarly pass them down if
+        # required
+        for var in dir(cls):
+            if var[0] == "_" or var in cls._NON_PARAMETERS or var in vars:
+                continue
+            val = getattr(cls, var)
+            # Exclude methods, properties and other classes
+            if isinstance(val, (MethodType, FunctionType, property, type)):
+                continue
+            setattr(cls, var, property(fget=lambda _, val=val: val, fset=set_cls_var))
+            vars.append(var)
+
         if passdown:
             self.flow._datastore.passdown_partial(parameter_ds, vars)
-        return vars
+        return param_only_vars
 
     def _init_data(self, run_id, join_type, input_paths):
         # We prefer to use the parallelized version to initialize datastores
diff --git a/test/core/metaflow_test/__init__.py b/test/core/metaflow_test/__init__.py
@@ -90,6 +90,7 @@ class MetaflowTest(object):
     PRIORITY = 999999999
     PARAMETERS = {}
     INCLUDE_FILES = {}
+    CLASS_VARS = {}
     HEADER = ""
 
     def check_results(self, flow, checker):
diff --git a/test/core/metaflow_test/formatter.py b/test/core/metaflow_test/formatter.py
@@ -88,6 +88,9 @@ def _flow_lines(self):
         yield 0, self.test.HEADER
         yield 0, "class %s(FlowSpec):" % self.flow_name
 
+        for var, val in self.test.CLASS_VARS.items():
+            yield 1, '%s = %s' % (var, val)
+
         for var, parameter in self.test.PARAMETERS.items():
             kwargs = ["%s=%s" % (k, v) for k, v in parameter.items()]
             yield 1, '%s = Parameter("%s", %s)' % (var, var, ",".join(kwargs))
diff --git a/test/core/tests/constants.py b/test/core/tests/constants.py
@@ -0,0 +1,54 @@
+from metaflow_test import MetaflowTest, ExpectationFailed, steps
+
+
+class ConstantsTest(MetaflowTest):
+    """
+    Test that an artifact defined in the first step
+    is available in all steps downstream.
+    """
+
+    PRIORITY = 0
+    CLASS_VARS = {'str_const': '"this is a constant"',
+                  'int_const': 123,
+                  'obj_const': '[]'}
+
+    PARAMETERS = {
+        "int_param": {"default": 456},
+        "str_param": {"default": "'foobar'"},
+    }
+
+    @steps(0, ["all"])
+    def step_all(self):
+        # make sure class attributes are available in all steps
+        # through joins etc
+        assert_equals('this is a constant', self.str_const)
+        assert_equals(123, self.int_const)
+        # obj_const is mutable. Not much that can be done about it
+        assert_equals([], self.obj_const)
+
+        assert_equals(456, self.int_param)
+        assert_equals('foobar', self.str_param)
+
+        # make sure class variables are not listed as parameters
+        from metaflow import current
+        assert_equals({'int_param', 'str_param'},
+                      set(current.parameter_names))
+
+        try:
+            self.int_param = 5
+        except AttributeError:
+            pass
+        else:
+            raise Exception("It shouldn't be possible to modify parameters")
+
+        try:
+            self.int_const = 122
+        except AttributeError:
+            pass
+        else:
+            raise Exception("It shouldn't be possible to modify constants")
+
+    def check_results(self, flow, checker):
+        for step in flow:
+            checker.assert_artifact(step.name, 'int_param', 456)
+            checker.assert_artifact(step.name, 'int_const', 123)
