fix: dynamic var spec serialization, JSON indexes, and StepMutator sequencing

metaflow/decorators.py

@@ -1,3 +1,4 @@
+import base64
 import importlib
 import json
 import re
@@ -15,6 +16,7 @@
 from .debug import debug
 from .dynamic_var import (
     DynamicVar,
+    _NO_DEFAULT,
     _contains_dynamic_var,
     has_dynamic_vars,
     resolve_dynamic_vars_from_store,
@@ -36,6 +38,9 @@
 from .metaflow_config import SPIN_ALLOWED_DECORATORS
 from metaflow._vendor import click
 
+_DYNAMIC_VAR_PREFIX = "__dynvar__:"
+_DYNAMIC_VAR_JSON_PREFIX = "json:"
+
 
 class BadStepDecoratorException(MetaflowException):
     headline = "Syntax error"
@@ -202,8 +207,10 @@ def extract_args_kwargs_from_decorator_spec(cls, deco_spec):
         for a in re.split(r""",(?=[\s\w]+=)""", deco_spec):
             name, val = a.split("=", 1)
             val_stripped = val.strip()
-            if val_stripped.startswith("__dynvar__:"):
-                val_parsed = DynamicVar(val_stripped[len("__dynvar__:") :])
+            if val_stripped.startswith(_DYNAMIC_VAR_PREFIX):
+                val_parsed = Decorator._decode_dynamic_var_sentinel(
+                    val_stripped[len(_DYNAMIC_VAR_PREFIX) :]
+                )
             else:
                 try:
                     val_parsed = json.loads(val_stripped.replace('\\"', '"'))
@@ -233,14 +240,47 @@ def parse_decorator_spec(cls, deco_spec):
         _, kwargs = cls.extract_args_kwargs_from_decorator_spec(deco_spec)
         return cls(attributes=kwargs)
 
+    @staticmethod
+    def _encode_dynamic_var_sentinel(value):
+        payload = {
+            "var_name": value.var_name,
+            "pertask": value.pertask,
+            "has_default": value.default is not _NO_DEFAULT,
+        }
+        if value.default is not _NO_DEFAULT:
+            payload["default"] = value.default
+        encoded = base64.urlsafe_b64encode(
+            json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
+        ).decode("ascii")
+        return "%s%s%s" % (_DYNAMIC_VAR_PREFIX, _DYNAMIC_VAR_JSON_PREFIX, encoded)
+
+    @staticmethod
+    def _decode_dynamic_var_sentinel(payload):
+        if payload.startswith(_DYNAMIC_VAR_JSON_PREFIX):
+            data = json.loads(
+                base64.urlsafe_b64decode(
+                    payload[len(_DYNAMIC_VAR_JSON_PREFIX) :].encode("ascii")
+                ).decode("utf-8")
+            )
+            if data.get("has_default"):
+                return DynamicVar(
+                    data["var_name"],
+                    pertask=bool(data.get("pertask", False)),
+                    default=data.get("default"),
+                )
+            return DynamicVar(
+                data["var_name"], pertask=bool(data.get("pertask", False))
+            )
+        return DynamicVar(payload)
+
     @staticmethod
     def _encode_dynamic_vars(val):
         """Recursively encode DynamicVar as short string sentinel for JSON serialization."""
 
         # DynamicVar as value or key
         def encode_dyn(v):
             if isinstance(v, DynamicVar):
-                return "__dynvar__:%s" % v.var_name
+                return Decorator._encode_dynamic_var_sentinel(v)
             return v
 
         if isinstance(val, dict):
@@ -266,8 +306,10 @@ def _decode_dynamic_vars(val):
         """Recursively decode __dynvar__ string sentinels back to DynamicVar,"""
 
         def try_decode(v):
-            if isinstance(v, str) and v.startswith("__dynvar__:"):
-                return DynamicVar(v[len("__dynvar__:") :])
+            if isinstance(v, str) and v.startswith(_DYNAMIC_VAR_PREFIX):
+                return Decorator._decode_dynamic_var_sentinel(
+                    v[len(_DYNAMIC_VAR_PREFIX) :]
+                )
             return v
 
         if isinstance(val, dict):
@@ -298,7 +340,9 @@ def make_decorator_spec(self):
             # we dump using JSON.
             for k, v in attrs.items():
                 if isinstance(v, DynamicVar):
-                    attr_list.append("%s=__dynvar__:%s" % (k, v.var_name))
+                    attr_list.append(
+                        "%s=%s" % (k, self._encode_dynamic_var_sentinel(v))
+                    )
                 elif isinstance(v, (int, float, str)):
                     attr_list.append("%s=%s" % (k, str(v)))
                 else:

metaflow/dynamic_var.py

@@ -216,8 +216,13 @@ def _resolve_val(val, source, split_index=None):
         if isinstance(collection, list):
             return collection[split_index]
         elif isinstance(collection, dict):
+            if split_index in collection:
+                return collection[split_index]
+            str_split_index = str(split_index)
+            if str_split_index in collection:
+                return collection[str_split_index]
             if val.default is not _NO_DEFAULT:
-                return collection.get(split_index, val.default)
+                return val.default
             return collection[split_index]
         else:
             from .exception import MetaflowException

metaflow/flowspec.py

@@ -449,7 +449,7 @@ def _process_config_decorators(cls, config_options, process_configs=True):
         # currently documented but follows the usual principle of finishing a callback
         # phase before starting the next one.
         for step in cls._steps:
-            for deco in step.decorators:
+            for deco in step.config_decorators:
                 if isinstance(deco, StepMutator):
                     deco.external_init()
                 else:
@@ -473,23 +473,6 @@ def _process_config_decorators(cls, config_options, process_configs=True):
                     )
                 )
 
-        for step in cls._steps:
-            for deco in step.config_decorators:
-                inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
-                debug.userconf_exec(
-                    "Evaluating step level decorator %s for %s (pre-mutate)"
-                    % (deco.__class__.__name__, step.name)
-                )
-                deco.pre_mutate(
-                    MutableStep(
-                        cls,
-                        step,
-                        pre_mutate=True,
-                        statically_defined=deco.statically_defined,
-                        inserted_by=inserted_by_value,
-                    )
-                )
-
         # Process parameters to allow them to also use config values easily
         for var, param in cls._get_parameters():
             if param.IS_CONFIG_PARAMETER:

test/unit/test_dynamic_var_decorator_spec.py

@@ -0,0 +1,75 @@
+from metaflow import var
+from metaflow.decorators import Decorator
+from metaflow.dynamic_var import _NO_DEFAULT, resolve_dynamic_vars
+from metaflow.plugins.catch_decorator import CatchDecorator
+from metaflow.plugins.environment_decorator import EnvironmentDecorator
+
+
+def _parse_decorator_attrs(decorator):
+    spec = decorator.make_decorator_spec()
+    _, attr_spec = spec.split(":", 1)
+    _, attrs = Decorator.extract_args_kwargs_from_decorator_spec(attr_spec)
+    return attrs
+
+
+def _assert_dynamic_var_spec(value, name, default):
+    assert value.var_name == name
+    assert value.pertask is True
+    assert value.default is not _NO_DEFAULT
+    assert value.default == default
+
+
+def test_legacy_dynamic_var_spec_still_parses():
+    _, attrs = Decorator.extract_args_kwargs_from_decorator_spec(
+        "var=__dynvar__:catch_name"
+    )
+
+    assert attrs["var"].var_name == "catch_name"
+    assert attrs["var"].pertask is False
+    assert attrs["var"].default is _NO_DEFAULT
+
+
+def test_top_level_dynamic_var_spec_preserves_pertask_and_default():
+    attrs = _parse_decorator_attrs(
+        CatchDecorator(
+            attributes={
+                "var": var("catch_names", pertask=True, default="caught_default"),
+                "print_exception": False,
+            },
+            statically_defined=True,
+        )
+    )
+
+    _assert_dynamic_var_spec(attrs["var"], "catch_names", "caught_default")
+
+
+def test_nested_dynamic_var_spec_preserves_pertask_and_default():
+    attrs = _parse_decorator_attrs(
+        EnvironmentDecorator(
+            attributes={
+                "vars": {
+                    "DYNAMIC_ENV_VALUE": var(
+                        "env_values", pertask=True, default="fallback"
+                    )
+                }
+            },
+            statically_defined=True,
+        )
+    )
+
+    _assert_dynamic_var_spec(
+        attrs["vars"]["DYNAMIC_ENV_VALUE"], "env_values", "fallback"
+    )
+
+
+def test_pertask_dict_resolution_accepts_json_stringified_indexes():
+    assert resolve_dynamic_vars(
+        {"var": var("catch_names", pertask=True, default="caught_default")},
+        {"catch_names": {"0": "caught_zero"}},
+        split_index=0,
+    ) == {"var": "caught_zero"}
+    assert resolve_dynamic_vars(
+        {"var": var("catch_names", pertask=True, default="caught_default")},
+        {"catch_names": {"0": "caught_zero"}},
+        split_index=1,
+    ) == {"var": "caught_default"}