PR Type

• Bug fix
• Core Runtime change (higher bar -- metaflow/decorators.py)

Summary

StepMutator.mutate() can again see and override late-attached platform decorators (e.g. @kubernetes). Before this fix, on deploy paths that attach the compute decorator late (Argo Workflows, --with conda/pypi, etc.) a StepMutator override of @kubernetes resources was silently dropped.

Issue

Fixes #3025

Reproduction

Runtime: argo-workflows (any deploy path that late-attaches @kubernetes)

A StepMutator that inspects decorator_specs for "kubernetes" and calls add_decorator(..., duplicates=OVERRIDE) to bump cpu/memory:

from metaflow import FlowSpec, StepMutator, step

class kubernetes_override(StepMutator):
    def init(self, *args, **kwargs):
        self.deco_kwargs = kwargs
    def mutate(self, mutable_step):
        for deco_name, _, _, attributes in mutable_step.decorator_specs:
            if deco_name == "kubernetes":
                attributes.update(self.deco_kwargs)
                mutable_step.add_decorator(
                    deco_type=deco_name, deco_kwargs=attributes,
                    duplicates=mutable_step.OVERRIDE,
                )

class MyFlow(FlowSpec):
    @kubernetes_override(cpu=2, memory=8192)
    @step
    def start(self):
        self.next(self.end)
    @step
    def end(self):
        pass

Commands to run:

python my_flow.py argo-workflows create --only-json

Where evidence shows up: the generated Argo workflow / the @kubernetes attributes on the start step.

start step @kubernetes: cpu=1, memory=4096   # defaults — the mutator never saw @kubernetes

start step @kubernetes: cpu=2, memory=8192   # mutator override is honored

The included unit test test/unit/test_late_attached_mutator.py encodes exactly this: it fails before the fix (AssertionError: assert '1' == '4') and passes after, with no cloud backend required — it drives the internal _init_step_decorators → _attach_decorators → _process_late_attached_decorator sequence directly using a MagicMock datastore.

Root Cause

The single caller of StepMutator.mutate() is _init_step_decorators (metaflow/decorators.py). PR #2719 moved _init_step_decorators() out of Argo's make_flow() and into cli.py start(), so the mutator pass now runs before platform decorators are late-attached:

1. _init_step_decorators() runs StepMutator.mutate() — @kubernetes is not attached yet
2. _attach_decorators() late-attaches @kubernetes to the steps
3. _process_late_attached_decorator() calls external_init() + step_init() on the new decorators but never re-runs mutate()

So a StepMutator iterating mutable_step.decorator_specs for "kubernetes" in step 1 finds nothing, and its add_decorator(..., OVERRIDE) is a no-op. The invariant violated: a StepMutator is supposed to observe every decorator on its step, including platform decorators attached at deploy time — true in 2.19.13, broken in 2.19.14.

Why This Fix Is Correct

_process_late_attached_decorator now re-runs step mutators on exactly the steps that received a late-attached decorator, then rebuilds the graph — mirroring the existing StepMutator handling in _init_step_decorators (same MutableStep(...) construction, same cls._init_graph() refresh). This restores the 2.19.13 ordering guarantee while keeping the 2.19.14 architecture. The fix is minimal (one function, +48 lines, no signature changes) and touches only the late-attach path.

Key correctness points:

• Scoped re-run: mutators are re-run only for late_attached_step_names (the steps that actually received a late-attached decorator), so mutators on unaffected steps are not invoked a second time.
• Idempotent external_init(): after add_decorator(OVERRIDE) may create replacement decorator instances, external_init() is called again; it is guarded by _ran_init, so already-initialized decorators are untouched.
• Graph rebuild only when needed: cls._init_graph() runs only if a mutator actually ran (mutators_ran), avoiding a needless rebuild on the common no-mutator path.

Failure Modes Considered

1. Non-idempotent / cross-step mutators. Re-running every mutator on every step could double-apply side effects or apply step A's override to step B. Mitigated by gating the re-run on late_attached_step_names and matching _init_step_decorators's per-step iteration; the test test_step_without_mutator_keeps_kubernetes_defaults asserts a mutator-less step keeps defaults (cpu=1, memory=4096) and that exactly one @kubernetes exists per step (no duplication from OVERRIDE).
2. No late attachment at all (the common case). If no decorator in deco_names is present, late_attached_step_names is empty, mutators_ran stays False, no mutator runs and no graph rebuild happens — behavior is identical to before. Covered by test_no_late_attachment_is_a_noop.
3. Backward compatibility / spin path. No public signature changed; the is_spin/skip_decorators handling in the subsequent step_init loop is untouched. The full test/unit suite passes unchanged.

Tests

• Unit tests added/updated — test/unit/test_late_attached_mutator.py (3 cases: override-applies, no-mutator-keeps-defaults, no-attach-noop)
• Reproduction script provided (the test is a runnable repro; fails before, passes after — no cloud backend needed)
• CI passes (tox -e unit locally; targeted decorator/graph/mutator suite: 207 passed, 0 regressions)

Non-Goals

FlowMutator.mutate() is not re-run here. A FlowMutator that inspects step-level decorators (via MutableFlow.steps) for a late-attached @kubernetes would have the symmetric blindness. I scoped this PR to the StepMutator case in the issue to keep the change minimal and reviewable. Happy to follow up with the FlowMutator re-run in a separate PR if you'd like it closed too — flagging it explicitly so it isn't silently missed.

AI Tool Usage

• AI tools were used (describe below)

• Tool: Claude Code.

• Used for: locating the regression (bisecting the ordering change to PR Remove a case where duplicate init functions were called (when deploy… #2719), drafting the fix and the unit test, and running the local test suite.

• Review: I reviewed and understand every line. The fix deliberately mirrors the existing StepMutator pass in _init_step_decorators. The test fails before the change (assert '1' == '4') and passes after, and the full test/unit suite passes with no regressions.