Modify the graph to allow for single step flows and custom names

Instead of enforcing a step named "start" and another named "end"
and a minimum of two steps, Metaflow now allows for single
step flows with any names. The "start" and "end" properties are
derived from the structure of the graph. We still require a single
entry point and a single exit point but they can be one and the same
and do not have to be named something specific.

Author: romain-intel

metaflow/client/core.py

@@ -2121,8 +2121,9 @@ def parent_steps(self) -> Iterator["Step"]:
             Parent step
         """
         graph_info = self.task["_graph_info"].data
+        start_step = graph_info.get("start_step", "start")
 
-        if self.id != "start":
+        if self.id != start_step:
             flow, run, _ = self.path_components
             for node_name, attributes in graph_info["steps"].items():
                 if self.id in attributes["next"]:
@@ -2139,8 +2140,9 @@ def child_steps(self) -> Iterator["Step"]:
             Child step
         """
         graph_info = self.task["_graph_info"].data
+        end_step = graph_info.get("end_step", "end")
 
-        if self.id != "end":
+        if self.id != end_step:
             flow, run, _ = self.path_components
             for next_step in graph_info["steps"][self.id]["next"]:
                 yield Step(f"{flow}/{run}/{next_step}", _namespace_check=False)
@@ -2153,7 +2155,7 @@ class Run(MetaflowObject):
     Attributes
     ----------
     data : MetaflowData
-        a shortcut to run['end'].task.data, i.e. data produced by this run.
+        A shortcut to the terminal step's task data produced by this run.
     successful : bool
         True if the run completed successfully.
     finished : bool
@@ -2165,7 +2167,7 @@ class Run(MetaflowObject):
     trigger : MetaflowTrigger
         Information about event(s) that triggered this run (if present). See `MetaflowTrigger`.
     end_task : Task
-        `Task` for the end step (if it is present already).
+        `Task` for the terminal step (if it is present already).
     """
 
     _NAME = "run"
@@ -2176,6 +2178,25 @@ def _iter_filter(self, x):
         # exclude _parameters step
         return x.id[0] != "_"
 
+    @property
+    def _graph_endpoints(self):
+        """
+        Returns (start_step_name, end_step_name) from _parameters metadata.
+
+        Falls back to ("start", "end") for runs that predate structural
+        inference (backward compatibility).
+        """
+        if not hasattr(self, "_cached_endpoints"):
+            start, end = "start", "end"
+            try:
+                params_meta = self["_parameters"].task.metadata_dict
+                start = params_meta.get("start_step", "start")
+                end = params_meta.get("end_step", "end")
+            except Exception:
+                pass
+            self._cached_endpoints = (start, end)
+        return self._cached_endpoints
+
     def steps(self, *tags: str) -> Iterator[Step]:
         """
         [Legacy function - do not use]
@@ -2298,17 +2319,18 @@ def finished_at(self) -> Optional[datetime]:
     @property
     def end_task(self) -> Optional[Task]:
         """
-        Returns the Task corresponding to the 'end' step.
+        Returns the Task corresponding to the terminal step.
 
-        This returns None if the end step does not yet exist.
+        This returns None if the terminal step does not yet exist.
 
         Returns
         -------
         Task, optional
-            The 'end' task
+            The terminal step's task
         """
         try:
-            end_step = self["end"]
+            _, end_step_name = self._graph_endpoints
+            end_step = self[end_step_name]
         except KeyError:
             return None
 
@@ -2481,8 +2503,9 @@ def trigger(self) -> Optional[Trigger]:
         Trigger, optional
             Container of triggering events
         """
-        if "start" in self and self["start"].task:
-            meta = self["start"].task.metadata_dict.get("execution-triggers")
+        start_step, _ = self._graph_endpoints
+        if start_step in self and self[start_step].task:
+            meta = self[start_step].task.metadata_dict.get("execution-triggers")
             if meta:
                 return Trigger(json.loads(meta))
         return None

metaflow/decorators.py

@@ -547,6 +547,19 @@ def _base_step_decorator(decotype, *args, **kwargs):
         func = args[0]
         if isinstance(func, (StepMutator, UserStepDecoratorBase)):
             func = func._my_step
+
+        # Step decorator applied to a class with a synthesized step method.
+        # Used by extensions that create a single synthetic @step on a
+        # FlowSpec subclass. _step_spec_step_name is set by the
+        # extension's metaclass.
+        if isinstance(func, type) and hasattr(func, "_step_spec_step_name"):
+            step_func = getattr(func, func._step_spec_step_name)
+            if hasattr(step_func, "is_step"):
+                step_func.decorators.append(
+                    decotype(attributes=kwargs, statically_defined=True)
+                )
+            return func
+
         if not hasattr(func, "is_step"):
             raise BadStepDecoratorException(decotype.name, func)
 
@@ -947,7 +960,11 @@ def step(
 
 
 def step(
-    f: Union[Callable[[FlowSpecDerived], None], Callable[[FlowSpecDerived, Any], None]],
+    f=None,
+    *,
+    start=False,
+    end=False,
+    node_info=None,
 ):
     """
     Marks a method in a FlowSpec as a Metaflow Step. Note that this
@@ -972,20 +989,41 @@ def foo(self):
 
     Parameters
     ----------
-    f : Union[Callable[[FlowSpecDerived], None], Callable[[FlowSpecDerived, Any], None]]
-        Function to make into a Metaflow Step
+    f : callable, optional
+        Function to make into a Metaflow Step. When using keyword arguments
+        (e.g. ``@step(start=True)``), this is ``None`` and a decorator
+        function is returned instead.
+    start : bool, default False
+        Mark this step as the start (entry) step of the flow.
+    end : bool, default False
+        Mark this step as the end (terminal) step of the flow.
+    node_info : dict, optional
+        Extra metadata to attach to this step's DAGNode. Extensions can use
+        this to store arbitrary information accessible via ``flow._graph``
+        (live references) and ``_graph_info`` (serialized via ``to_pod``).
 
     Returns
     -------
-    Union[Callable[[FlowSpecDerived, StepFlag], None], Callable[[FlowSpecDerived, Any, StepFlag], None]]
-        Function that is a Metaflow Step
+    callable
+        The decorated function, or a decorator if keyword arguments were used.
     """
-    f.is_step = True
-    f.decorators = []
-    f.config_decorators = []
-    f.wrappers = []
-    f.name = f.__name__
-    return f
+
+    def _apply(func):
+        func.is_step = True
+        func.decorators = []
+        func.config_decorators = []
+        func.wrappers = []
+        func.name = func.__name__
+        func.is_start_step = start
+        func.is_end_step = end
+        func.node_info = node_info or {}
+        return func
+
+    if f is not None:
+        # Called as @step (no parens)
+        return _apply(f)
+    # Called as @step(start=True) etc.
+    return _apply
 
 
 def _import_plugin_decorators(globals_dict):

metaflow/flowspec.py

@@ -164,11 +164,21 @@ def _merge_value(self, inherited_value, self_value):
 
 
 class FlowSpecMeta(type):
+    _registry = {}  # {class_name: class} for all FlowSpec/StepSpec subclasses
+
+    # Names of base classes that should NOT be registered or have their
+    # graph/attrs initialized.  Subclasses of FlowSpecMeta (like StepSpecMeta)
+    # can extend this set.
+    _base_class_names = frozenset({"FlowSpec"})
+
     def __init__(cls, name, bases, attrs):
         super().__init__(name, bases, attrs)
-        if name == "FlowSpec":
+        # Check against the metaclass's own _base_class_names — this
+        # allows StepSpecMeta to add "StepSpec" to the set.
+        if name in type(cls)._base_class_names:
             return
 
+        type(cls)._registry[name] = cls
         cls._init_attrs()
 
     def _init_attrs(cls):
@@ -516,6 +526,8 @@ def _set_constants(self, graph, kwargs, config_options):
 
         graph_info = {
             "file": os.path.basename(os.path.abspath(sys.argv[0])),
+            "start_step": graph.start_step,
+            "end_step": graph.end_step,
             "parameters": parameters_info,
             "constants": constants_info,
             "steps": steps_info,

metaflow/graph.py

@@ -48,7 +48,17 @@ def deindent_docstring(doc):
 
 class DAGNode(object):
     def __init__(
-        self, func_ast, decos, wrappers, config_decorators, doc, source_file, lineno
+        self,
+        func_ast,
+        decos,
+        wrappers,
+        config_decorators,
+        doc,
+        source_file,
+        lineno,
+        is_start_step=False,
+        is_end_step=False,
+        node_info=None,
     ):
         self.name = func_ast.name
         self.source_file = source_file
@@ -60,6 +70,12 @@ def __init__(
         self.config_decorators = config_decorators
         self.doc = deindent_docstring(doc)
         self.parallel_step = any(getattr(deco, "IS_PARALLEL", False) for deco in decos)
+        # Explicit start/end annotations from @step(start=True) / @step(end=True)
+        self.is_start_step = is_start_step
+        self.is_end_step = is_end_step
+        # Generic metadata dict for extensions to attach extra info to this node.
+        # Serialized to _graph_info via to_pod; live references accessible via flow._graph.
+        self.node_info = node_info or {}
 
         # these attributes are populated by _parse
         self.tail_next_lineno = 0
@@ -140,10 +156,9 @@ def _parse(self, func_ast, lineno):
         self.num_args = len(func_ast.args.args)
         tail = func_ast.body[-1]
 
-        # end doesn't need a transition
-        if self.name == "end":
-            # TYPE: end
-            self.type = "end"
+        # Note: type assignment for start/end steps is handled by
+        # FlowGraph._identify_start_end() based on graph structure,
+        # not by name.
 
         # ensure that the tail an expression
         if not isinstance(tail, ast.Expr):
@@ -212,10 +227,10 @@ def _parse(self, func_ast, lineno):
                     self.type = "split"
                     self.invalid_tail_next = False
                 elif len(self.out_funcs) == 1:
-                    # TYPE: linear
-                    if self.name == "start":
-                        self.type = "start"
-                    elif self.num_args > 1:
+                    # TYPE: linear (or join)
+                    # Note: "start" type is assigned later by
+                    # FlowGraph._identify_start_end() based on structure.
+                    if self.num_args > 1:
                         self.type = "join"
                     else:
                         self.type = "linear"
@@ -259,9 +274,65 @@ def __init__(self, flow):
         self.doc = deindent_docstring(flow.__doc__)
         # nodes sorted in topological order.
         self.sorted_nodes = []
+        self._identify_start_end()
         self._traverse_graph()
         self._postprocess()
 
+    def _identify_start_end(self):
+        """
+        Determine the start and end steps.
+
+        Uses explicit ``@step(start=True)`` / ``@step(end=True)`` annotations
+        if present.  Falls back to looking for steps named ``"start"`` /
+        ``"end"`` for backward compatibility.
+
+        Sets ``self.start_step`` and ``self.end_step`` to step name strings,
+        or ``None`` if the graph is malformed (validated later by lint).
+        Also assigns the ``"start"`` and ``"end"`` node types.
+        """
+        # 1. Look for explicit annotations
+        annotated_start = [
+            name
+            for name, node in self.nodes.items()
+            if node.is_start_step and not name.startswith("_")
+        ]
+        annotated_end = [
+            name
+            for name, node in self.nodes.items()
+            if node.is_end_step and not name.startswith("_")
+        ]
+
+        # 2. Determine start step (annotation first, then name fallback)
+        if len(annotated_start) == 1:
+            self.start_step = annotated_start[0]
+        elif len(annotated_start) == 0:
+            self.start_step = "start" if "start" in self.nodes else None
+        else:
+            self.start_step = None  # Multiple annotated — lint will catch
+
+        # 3. Determine end step (annotation first, then name fallback)
+        if len(annotated_end) == 1:
+            self.end_step = annotated_end[0]
+        elif len(annotated_end) == 0:
+            self.end_step = "end" if "end" in self.nodes else None
+        else:
+            self.end_step = None  # Multiple annotated — lint will catch
+
+        # 4. Assign types based on identified start/end.
+        # Only upgrade "linear" → "start" for the entry point; do NOT override
+        # "split", "foreach", etc. since those types are needed for
+        # split/join balance checking.
+        if self.start_step and self.start_step == self.end_step:
+            # Single-step flow: terminal node that is also the entry point
+            self.nodes[self.start_step].type = "end"
+        else:
+            if self.start_step:
+                node = self.nodes[self.start_step]
+                if node.type in (None, "linear"):
+                    node.type = "start"
+            if self.end_step:
+                self.nodes[self.end_step].type = "end"
+
     def _create_nodes(self, flow):
         nodes = {}
         for element in dir(flow):
@@ -281,6 +352,9 @@ def _create_nodes(self, flow):
                     func.__doc__,
                     source_file,
                     lineno,
+                    is_start_step=getattr(func, "is_start_step", False),
+                    is_end_step=getattr(func, "is_end_step", False),
+                    node_info=getattr(func, "node_info", None),
                 )
                 nodes[element] = node
         return nodes
@@ -338,8 +412,8 @@ def traverse(node, seen, split_parents, split_branches):
                             split_branches + ([n] if add_split_branch else []),
                         )
 
-        if "start" in self:
-            traverse(self["start"], [], [], [])
+        if self.start_step and self.start_step in self:
+            traverse(self[self.start_step], [], [], [])
 
         # fix the order of in_funcs
         for node in self.nodes.values():
@@ -445,6 +519,7 @@ def node_to_dict(name, node):
                     for deco in chain(node.wrappers, node.config_decorators)
                 ],
                 "next": node.out_funcs,
+                "node_info": to_pod(node.node_info),
             }
             if d["type"] == "split-foreach":
                 d["foreach_artifact"] = node.foreach_param
@@ -493,9 +568,15 @@ def populate_block(start_name, end_name):
                         break
             return resulting_list
 
-        graph_structure = populate_block("start", "end")
+        if self.start_step == self.end_step:
+            # Single-step flow
+            graph_structure = []
+        else:
+            graph_structure = populate_block(self.start_step, self.end_step)
 
-        steps_info["end"] = node_to_dict("end", self.nodes["end"])
-        graph_structure.append("end")
+        steps_info[self.end_step] = node_to_dict(
+            self.end_step, self.nodes[self.end_step]
+        )
+        graph_structure.append(self.end_step)
 
         return steps_info, graph_structure