DAG data leaks raw StepMutator constructor args, breaking Metaflow UI when the argument's class isn't importable

[sertansenturk-deliveroo]

Summary

Metaflow UI fails to display the DAG for flows using a StepMutator whose constructor takes a non-primitive argument. graph.py's DAG serialization sanitizes classic step decorator attributes with to_pod(), but not StepMutator/config-decorator attributes, so a raw object gets pickled into _graph_info. Any consumer that can't import the object's class (particularly the Metaflow UI backend) fails to unpickle it.

Expected: the DAG renders regardless of which third-party class a StepMutator's constructor argument came from, the same way it already does for classic step decorators.

Actual: the DAG tab fails with a stack trace ending in a ModuleNotFoundError for the argument's class, whenever the consumer reading _graph_info doesn't have that class importable.

Where

In node_to_dict():

Classic decorator attributes go through to_pod():

metaflow/metaflow/graph.py

Lines 615 to 624 in 4fce948

	"decorators": [
	{
	"name": deco.name,
	"attributes": to_pod(deco.attributes),
	"statically_defined": deco.statically_defined,
	"inserted_by": deco.inserted_by,
	}
	for deco in node.decorators
	if not deco.name.startswith("_")
	]

StepMutator/config-decorator attributes do not:

metaflow/metaflow/graph.py

Lines 625 to 633 in 4fce948

	+ [
	{
	"name": deco.decorator_name,
	"attributes": {"_args": deco._args, **deco._kwargs},
	"statically_defined": deco.statically_defined,
	"inserted_by": deco.inserted_by,
	}
	for deco in chain(node.wrappers, node.config_decorators)
	],

{
    "name": deco.decorator_name,
    "attributes": {"_args": deco._args, **deco._kwargs},  # <-- no to_pod()
    "statically_defined": deco.statically_defined,
    "inserted_by": deco.inserted_by,
}
for deco in chain(node.wrappers, node.config_decorators)

How to Reproduce

payload_module.py:

class Payload:
    def __init__(self, x):
        self.x = x

repro_flow.py:

from metaflow import FlowSpec, step
from metaflow.user_decorators.user_step_decorator import StepMutator
from metaflow.user_decorators.mutable_step import MutableStep

from payload_module import Payload


class MyMutator(StepMutator):
    def init(self, payload=None, **kwargs):
        self._payload = payload

    def mutate(self, mutable_step: MutableStep) -> None:
        pass


class ReproFlow(FlowSpec):
    @MyMutator(payload=Payload(1))
    @step
    def start(self):
        self.next(self.end)

    @step
    def end(self):
        pass


if __name__ == "__main__":
    ReproFlow()

1. python repro_flow.py run
2. Read the run's _graph_info artifact. steps["start"]["decorators"] contains a raw Payload instance under attributes["payload"]. Verified:

   from metaflow import namespace, Run
   namespace(None)
   run = Run("ReproFlow/<run-id>")
   gi = run["start"].task["_graph_info"].data
   print(gi["steps"]["start"]["decorators"])
   # [{'name': '__main__.MyMutator', 'attributes': {'_args': [], 'payload': <payload_module.Payload object at ...>}, ...}]

3. Move payload_module.py out of the way and read _graph_info again from the same directory (so the local datastore still resolves, only the module becomes unimportable). It raises:

   ModuleNotFoundError: No module named 'payload_module'
   

   This is the same failure mode the Metaflow UI backend hits: it has no reason to have any flow-specific module installed, yet ends up needing it just to unpickle _graph_info.

Impact

Any third-party StepMutator/config decorator whose constructor accepts a non-primitive argument breaks the Metaflow UI's DAG tab for every run of every flow using it, since the UI backend has no reason to have that decorator's package installed. This has been the case since StepMutator/config_decorators was introduced in 2.16.0 (#2463), and the same code path is unchanged on the current master (commit 4fce948c).

Suggested fix

Apply to_pod() to the StepMutator/config-decorator attributes the same way it's applied to classic decorator attributes, in the second branch of node_to_dict()'s decorators list.

Environment

• Metaflow 2.19.35, reproduced locally (no cloud/k8s involved; this is a pure DAG-serialization bug).
• Same code path confirmed present on master at commit 4fce948c811aa4c9b9958f77367d527cf2695921.
• macOS 15 (Darwin 25.6.0, arm64), Python 3.11.15.