Skip to content

Allow UserStepDecorators to also be StepMutators#3079

Merged
romain-intel merged 25 commits into
masterfrom
romain/worktree/hpo
Apr 24, 2026
Merged

Allow UserStepDecorators to also be StepMutators#3079
romain-intel merged 25 commits into
masterfrom
romain/worktree/hpo

Conversation

@romain-intel
Copy link
Copy Markdown
Contributor

@romain-intel romain-intel commented Mar 31, 2026

Summary

Allow a class to inherit from both UserStepDecorator and StepMutator at the same time, so a single decorator can contribute both a mutate()/pre_mutate() pass over the flow/step (the StepMutator role) and pre_step/post_step wrapping hooks (the UserStepDecorator role). The instance is registered on both step.wrappers and step.config_decorators via a new MRO-walking _all_step_fields() helper, with an _ran_init idempotency guard so external_init() is not called twice.

In the process we also tightened a number of nearby corners in the user-decorator / flow-decorator plumbing and added several public helpers that made the dual-inheritance pattern writable at all.

PR Type

  • Bug fix
  • New feature
  • Core Runtime change
  • Docs / tooling
  • Refactoring

New functionality

Dual inheritanceclass foo(UserStepDecorator, StepMutator) is now a first-class pattern:

  • Registration is unified via _all_step_fields() which walks the MRO and returns every _step_field declared on ancestor classes (deterministic MRO order).
  • _ran_init guard on external_init() ensures init() is only invoked once per instance regardless of how many fields the instance lives on.
  • _init_step_decorators now calls external_init on step.wrappers too (previously only step.decorators); _process_config_decorators covers config_decorators. This fixes a pre-existing gap where a pure UserStepDecorator wrapper's init() could be skipped.

New MutableStep / MutableFlow helpers (to support writing dual-inherit decorators without reaching into internal lists):

  • MutableFlow.get_step(name)
  • MutableFlow.has_decorator(name) / MutableStep.has_decorator(name)
  • MutableFlow.get_decorator_kwargs(...) / MutableStep.get_decorator_kwargs(...)
  • MutableStep.get_all_decorator_kwargs(...)
  • MutableStep.replace_decorator(...)
  • MutableStep.decorator_specs dedupes instances that live on multiple step fields.

New StepDecorator hooks:

  • check_step_conflicts(decorators) — called during _init_step_decorators so a decorator can reject incompatible neighbors early.
  • find_decorator_attrs(...) convenience for decorator-to-decorator introspection.

Task-runtime plumbing:

  • task.py post-step sentinel changed from False to None, with explicit is not False guard so post_step returning (None, False) is a no-op instead of tripping "Invalid value passed to self.next".
  • cli.py no longer keeps the old manual _init() / external_init() loops for flow mutators — it's all driven by _process_config_decorators and _init_step_decorators now.
  • The standalone decorators._init() helper was removed; its responsibilities were absorbed into _init_flow_decorators and _init_step_decorators.

Bug fixes (most flagged by Greptile during review)

  • add_decorator on FlowMutator — fixed AttributeError: deco_type.name when add_decorator(MyFlowMutator) is called from mutate(); the error path now uses getattr(deco_type, "name", None) or getattr(deco_type, "decorator_name", deco_type) so the real MetaflowException is raised.
  • remove_decorator("typo")get_decorator_by_name returns None for unregistered names; no longer passes that through to issubclass(None, ...) and raises TypeError.
  • FlowMutator registration_do_add now calls external_init() before appending to the decorator lists, so a raising external_init doesn't leave a half-initialized mutator in _self_data / merged mutators.
  • OVERRIDE + merged list — switched from list.remove() (which shifts indices under the outer pre_mutate iteration) to slice-assignment merged[:] = [...]. No more silently-skipped elements.
  • cached_skip_decorators — removed a name-based cache that over-skipped allow_multiple decorators (e.g. multiple @cards on different steps). Matches master's per-decorator behavior.
  • remove_decorator(name) do_all=True — the _pre_mutate + isinstance(StepMutator) guard is now enforced on both the specific-match and do_all paths; previously do_all could strip a StepMutator from mutate() in violation of the documented contract.
  • add_decorator IGNORE path — the UserStepDecoratorBase branches now return None when add_or_raise swallowed the registration (duplicate + IGNORE, or non-statically-defined ERROR), matching the StepDecorator path and the documented contract. Previously the throwaway instance was returned as if it had been added.

Tests

  • New test/unit/mutators/ suite with dedicated flows + tests for dual inheritance, dynamic flow-mutator addition, string-form mutator addition, add_decorator return value, and regression coverage for the fixes above (test_remove_decorator_guard.py, test_post_step_none_false.py).
  • Dual inheritance test exercises all four lifecycle methods: pre_mutate, mutate, pre_step, post_step.
  • Spin conftest registers --use-latest so the fixture factory works when the suite is driven by a batch runner that doesn't inherit the option from a parent conftest.

@greptile-apps
Copy link
Copy Markdown
Contributor

greptile-apps Bot commented Mar 31, 2026

Greptile Summary

This PR enables a class to simultaneously inherit from both UserStepDecorator and StepMutator, running pre_mutate/mutate at config time and pre_step/post_step at task time. It also extends MutableFlow.add_decorator to accept FlowMutator subclasses dynamically and adds has_decorator/get_decorator_specs helpers to both MutableStep and MutableFlow. All previously-reported review concerns (task.py False sentinel, add_or_raise OVERRIDE duplicate, decorator_specs double-yield, external_init not reached for wrappers, do_all=True bypassing the pre-mutate guard, TypeError on unknown decorator name in remove_decorator, and AttributeError on FlowMutator.name) are addressed in this revision.

Confidence Score: 4/5

Safe to merge; all previously-flagged P0/P1 issues are resolved in this revision

The remaining complexity is in the dual-inheritance lifecycle (five call sites for external_init across three files) which is correctly guarded by _ran_init, but the non-trivial interaction surface warrants a 4 rather than 5

metaflow/user_decorators/user_step_decorator.py (add_or_raise / _all_step_fields) and metaflow/user_decorators/mutable_step.py (add_decorator string-path)

Important Files Changed

Filename Overview
metaflow/task.py Fixes sentinel from False→None and adds or None coercion plus is not False guard so post_step returning (None, False) is correctly a no-op instead of raising RuntimeError
metaflow/user_decorators/user_step_decorator.py Core change: adds _all_step_fields(), _ran_init guard, and multi-field registration to support dual UserStepDecorator+StepMutator inheritance; OVERRIDE cleanup now removes from all fields
metaflow/user_decorators/mutable_step.py Adds seen_decorators deduplication in decorator_specs, has_decorator/get_decorator_specs helpers, external_init on newly-added decorators, and correct do_all+pre_mutate guard in remove_decorator
metaflow/user_decorators/mutable_flow.py Adds FlowMutator addition/duplicate-handling via add_decorator, has_decorator/get_decorator_specs, decorator_specs extended to include FlowMutators, and guard against removing FlowMutators during iteration
metaflow/decorators.py Removes _init() and consolidates external_init calls: flow decorators in _init_flow_decorators, step.decorators and step.wrappers in _init_step_decorators, config_decorators in _process_config_decorators
metaflow/flowspec.py Splits _process_config_decorators: calls external_init on FLOW_MUTATORS first, then external_init on config_decorators before pre_mutate, separating init and mutation phases correctly
metaflow/user_decorators/user_flow_decorator.py Adds _ran_init guard to FlowMutator.external_init, adds parse_decorator_spec and extract_args_kwargs_from_decorator_spec class methods for string-based addition

Reviews (17): Last reviewed commit: "add_decorator IGNORE path: return None i..." | Re-trigger Greptile

Comment thread metaflow/task.py Outdated
Comment thread metaflow/user_decorators/user_step_decorator.py Outdated
@romain-intel romain-intel force-pushed the romain/worktree/hpo branch 2 times, most recently from 12f3a80 to 86e5908 Compare April 14, 2026 19:57
Comment thread metaflow/decorators.py
def init(self, marker="default"):
self._marker = marker

def mutate(self, mutable_step):
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could also cover pre_mutate, but validating one or the other should be enough of a confirmation as there's no cross-talk between the two?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I could add both but yes, it should not have cross talk.

Comment thread metaflow/decorators.py
return False


def _init(flow, only_non_static=False):
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so if an extension has been relying on attaching decorators, and needing to init these, instead of decorators._init(), will deco.external_init() now serve the same purpose?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So yes, this function just called all the external_init functions. I now spread out the external_init in different parts so that it is more accurate about when to call them.

# These decospecs are the ones from run/resume/spin PLUS the ones from the
# environment (for example the @conda)
decorators._attach_decorators(obj.flow, all_decospecs)
decorators._init(obj.flow)
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is an init not required anymore?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's still there but the calls are spread out.

Comment thread metaflow/flowspec.py
"Evaluating step level decorator %s for %s (pre-mutate)"
% (deco.__class__.__name__, step.name)
)
deco.pre_mutate(
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just a review note: this is fine, because we check in the previous loop that all decos in the config_decorators are StepMutators.

What is a bit unclear is why there was a need to split these to two loops. Is calling external_init() on every config deco before mutating a requirement?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was trying to keep with the other hooks where all hooks are called first and then all the second one.

romain-intel and others added 16 commits April 23, 2026 11:10
Also minor changes in error messages
 - Clean up "hooks" for mutators
 - Allow flow mutators to add other flow mutators
 - Fix an issue with step mutators adding step mutators
 - Return the decorator added with add_decorator
Allows looking up a step's mutable proxy by name without iterating
through all steps. Raises MetaflowException if the step doesn't exist.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- has_decorator(name): check if a step has a decorator by name
- get_decorator_kwargs(name): get the kwargs of the first matching decorator

Enables cleaner decorator inspection without manual tuple iteration.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Merge-and-replace a decorator's kwargs in one call, with optional
reject_if for declaring incompatible options.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Returns all decorator kwargs for a given name, enabling grouped queries
without manual iteration.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Called automatically during step_init to validate this decorator
against all others on the same step. Subclasses override to declare
incompatible combinations.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Runtime equivalent of MutableStep.get_decorator_kwargs — finds a
decorator's attributes on an instantiated step, with optional
exclude filter for skipping specific attribute values.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Same query API as MutableStep, but for flow-level decorators.
Eliminates manual iteration of flow.decorator_specs.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The error message referenced deco_type.name which doesn't exist on
FlowMutator classes (they use decorator_name). This caused an opaque
AttributeError instead of the intended MetaflowException when calling
add_decorator(SomeFlowMutator) from mutate().

Use getattr fallback chain: .name -> .decorator_name -> repr.
The dual_deco test now exercises all four lifecycle methods for a
class inheriting from both UserStepDecorator and StepMutator:
pre_mutate, mutate, pre_step, and post_step.
get_decorator_by_name returns None for unregistered names, which
caused issubclass(None, ...) to raise TypeError. Add a None guard
so unregistered names fall through to the StepDecorator name-based
removal path gracefully.
- CRITICAL-1: Call external_init() before appending to lists so a
  failure doesn't leave a half-initialized mutator in the decorator
  registries.

- MAJOR-1: OVERRIDE mode now uses slice-assignment (merged[:] = ...)
  instead of list.remove() to avoid skipping elements in the outer
  pre_mutate iteration loop that holds a reference to the same list.
The name-based caching of skipped decorators would cause all instances
of an allow_multiple decorator (e.g. @card) to be skipped if any single
instance was skipped for spin mode. Revert to per-decorator checking,
matching the behavior on master.
The flow_fixture factory calls request.config.getoption("--use-latest")
but the option was never registered via pytest_addoption in this
conftest. It worked when run from a parent conftest that registered
the option, but broke when the spin tests were invoked from a batch
runner that doesn't pull in that parent.

Register the option locally and pass default=False to getoption() for
belt-and-suspenders safety.
saikonen
saikonen previously approved these changes Apr 24, 2026
Copy link
Copy Markdown
Collaborator

@saikonen saikonen left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

extension changes verified to work, should be good to go from that side.

Previously `remove_decorator(name)` called without deco_args/deco_kwargs
(do_all=True) would skip the `_pre_mutate + isinstance(deco, StepMutator)`
check via its early `continue`, letting a user strip a StepMutator from
`mutate()` in violation of the documented contract.

Restructure the match/guard so both the do_all path and the specific-match
path go through the same StepMutator-from-mutate check.

Add a unit test covering the guard in both directions.
When a UserStepDecorator's post_step returned (None, False), task.py
would unpack the tuple (raised_exception=None, fake_next_call_args=False),
then reach the "Invalid value passed to self.next" RuntimeError branch
because False matches neither True, {}, nor had_raised_exception.

False is the natural "don't override self.next" return shape, semantically
identical to None. Accept both at the early guard so (None, False) is a
no-op as users reasonably expect.

Add a regression flow + tests exercising this return shape.
When `MutableStep.add_decorator` handles a UserStepDecoratorBase and
`add_or_raise` suppresses the registration (duplicates=IGNORE, or
non-statically-defined ERROR), the decorator was never actually added
to the step. But the outer `return step_deco` / `return d` was still
handing back the throwaway instance, contradicting the documented
"None if not added due to duplicate handling" contract and diverging
from the StepDecorator path which correctly returns None via the
added_deco tracker.

Make `add_or_raise` return `self` when the decorator was registered
and `None` when the call was swallowed. Propagate that signal at both
UserStepDecoratorBase call sites (string-form and class-form) so the
public return value matches the contract and the StepDecorator path.

Also fix _all_step_fields non-determinism: it used a set for dedup and
returned `list(set)`, giving non-deterministic ordering. Switched to a
parallel ordered list + dedup set so callers see a deterministic MRO
order (matters for side-effecting traversals like OVERRIDE cleanup).
@romain-intel romain-intel merged commit fc105a3 into master Apr 24, 2026
45 checks passed
@romain-intel romain-intel deleted the romain/worktree/hpo branch April 24, 2026 19:27
agsaru pushed a commit to agsaru/metaflow that referenced this pull request May 1, 2026
## Summary

Allow a class to inherit from both `UserStepDecorator` and `StepMutator`
at the same time, so a single decorator can contribute both a
`mutate()`/`pre_mutate()` pass over the flow/step (the `StepMutator`
role) and `pre_step`/`post_step` wrapping hooks (the `UserStepDecorator`
role). The instance is registered on both `step.wrappers` and
`step.config_decorators` via a new MRO-walking `_all_step_fields()`
helper, with an `_ran_init` idempotency guard so `external_init()` is
not called twice.

In the process we also tightened a number of nearby corners in the
user-decorator / flow-decorator plumbing and added several public
helpers that made the dual-inheritance pattern writable at all.

## PR Type
- [x] Bug fix
- [x] New feature
- [ ] Core Runtime change
- [ ] Docs / tooling
- [ ] Refactoring

## New functionality

**Dual inheritance** — `class foo(UserStepDecorator, StepMutator)` is
now a first-class pattern:
- Registration is unified via `_all_step_fields()` which walks the MRO
and returns every `_step_field` declared on ancestor classes
(deterministic MRO order).
- `_ran_init` guard on `external_init()` ensures `init()` is only
invoked once per instance regardless of how many fields the instance
lives on.
- `_init_step_decorators` now calls `external_init` on `step.wrappers`
too (previously only `step.decorators`); `_process_config_decorators`
covers `config_decorators`. This fixes a pre-existing gap where a pure
`UserStepDecorator` wrapper's `init()` could be skipped.

**New MutableStep / MutableFlow helpers** (to support writing
dual-inherit decorators without reaching into internal lists):
- `MutableFlow.get_step(name)`
- `MutableFlow.has_decorator(name)` / `MutableStep.has_decorator(name)`
- `MutableFlow.get_decorator_kwargs(...)` /
`MutableStep.get_decorator_kwargs(...)`
- `MutableStep.get_all_decorator_kwargs(...)`
- `MutableStep.replace_decorator(...)`
- `MutableStep.decorator_specs` dedupes instances that live on multiple
step fields.

**New StepDecorator hooks:**
- `check_step_conflicts(decorators)` — called during
`_init_step_decorators` so a decorator can reject incompatible neighbors
early.
- `find_decorator_attrs(...)` convenience for decorator-to-decorator
introspection.

**Task-runtime plumbing:**
- `task.py` post-step sentinel changed from `False` to `None`, with
explicit `is not False` guard so `post_step` returning `(None, False)`
is a no-op instead of tripping `"Invalid value passed to self.next"`.
- `cli.py` no longer keeps the old manual `_init()` / `external_init()`
loops for flow mutators — it's all driven by
`_process_config_decorators` and `_init_step_decorators` now.
- The standalone `decorators._init()` helper was removed; its
responsibilities were absorbed into `_init_flow_decorators` and
`_init_step_decorators`.

## Bug fixes (most flagged by Greptile during review)

- **`add_decorator` on FlowMutator** — fixed `AttributeError:
deco_type.name` when `add_decorator(MyFlowMutator)` is called from
`mutate()`; the error path now uses `getattr(deco_type, "name", None) or
getattr(deco_type, "decorator_name", deco_type)` so the real
`MetaflowException` is raised.
- **`remove_decorator("typo")`** — `get_decorator_by_name` returns
`None` for unregistered names; no longer passes that through to
`issubclass(None, ...)` and raises `TypeError`.
- **FlowMutator registration** — `_do_add` now calls `external_init()`
**before** appending to the decorator lists, so a raising
`external_init` doesn't leave a half-initialized mutator in `_self_data`
/ merged mutators.
- **OVERRIDE + merged list** — switched from `list.remove()` (which
shifts indices under the outer `pre_mutate` iteration) to
slice-assignment `merged[:] = [...]`. No more silently-skipped elements.
- **`cached_skip_decorators`** — removed a name-based cache that
over-skipped `allow_multiple` decorators (e.g. multiple `@card`s on
different steps). Matches master's per-decorator behavior.
- **`remove_decorator(name)` `do_all=True`** — the `_pre_mutate +
isinstance(StepMutator)` guard is now enforced on both the
specific-match and `do_all` paths; previously `do_all` could strip a
`StepMutator` from `mutate()` in violation of the documented contract.
- **`add_decorator` IGNORE path** — the `UserStepDecoratorBase` branches
now return `None` when `add_or_raise` swallowed the registration
(duplicate + IGNORE, or non-statically-defined ERROR), matching the
`StepDecorator` path and the documented contract. Previously the
throwaway instance was returned as if it had been added.

## Tests

- New `test/unit/mutators/` suite with dedicated flows + tests for dual
inheritance, dynamic flow-mutator addition, string-form mutator
addition, `add_decorator` return value, and regression coverage for the
fixes above (`test_remove_decorator_guard.py`,
`test_post_step_none_false.py`).
- Dual inheritance test exercises all four lifecycle methods:
`pre_mutate`, `mutate`, `pre_step`, `post_step`.
- Spin conftest registers `--use-latest` so the fixture factory works
when the suite is driven by a batch runner that doesn't inherit the
option from a parent conftest.

---------

Co-authored-by: Nissan Pow <npow@netflix.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants