Address comments

Author: romain-intel

metaflow/decorators.py

@@ -957,7 +957,10 @@ def step(
 
 
 def step(
-    f: Union[Callable[[FlowSpecDerived], None], Callable[[FlowSpecDerived, Any], None]],
+    f=None,
+    *,
+    start=False,
+    end=False,
 ):
     """
     Marks a method in a FlowSpec as a Metaflow Step. Note that this
@@ -982,20 +985,36 @@ 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.
 
     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
+        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/graph.py

@@ -48,7 +48,16 @@ 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,
     ):
         self.name = func_ast.name
         self.source_file = source_file
@@ -60,6 +69,9 @@ 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
 
         # these attributes are populated by _parse
         self.tail_next_lineno = 0
@@ -264,41 +276,45 @@ def __init__(self, flow):
 
     def _identify_start_end(self):
         """
-        Determine the start and end steps from graph structure.
+        Determine the start and end steps.
 
-        Start step: the unique node with zero in-degree (no other node
-        transitions to it). End step: the unique node with zero out-degree
-        (it has no self.next() transitions). Internal steps (names starting
-        with '_') are excluded.
+        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 the step name strings,
-        or None if the graph is malformed (validated later by lint).
-        Also assigns the "start" and "end" node types based on structure.
+        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.
         """
-        # Compute in-degree from out_funcs (already set by _parse)
-        in_degree = {name: 0 for name in self.nodes}
-        for node in self.nodes.values():
-            for target in node.out_funcs:
-                if target in in_degree:
-                    in_degree[target] += 1
-
-        # Start = zero in-degree (exclude internal steps starting with _)
-        candidates_start = [
+        # 1. Look for explicit annotations
+        annotated_start = [
             name
-            for name, deg in in_degree.items()
-            if deg == 0 and not name.startswith("_")
+            for name, node in self.nodes.items()
+            if node.is_start_step and not name.startswith("_")
         ]
-        # End = zero out-degree (exclude internal steps starting with _)
-        candidates_end = [
+        annotated_end = [
             name
-            for name in self.nodes
-            if not self.nodes[name].out_funcs and not name.startswith("_")
+            for name, node in self.nodes.items()
+            if node.is_end_step and not name.startswith("_")
         ]
 
-        self.start_step = candidates_start[0] if len(candidates_start) == 1 else None
-        self.end_step = candidates_end[0] if len(candidates_end) == 1 else None
+        # 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
 
-        # Assign types based on structure.
+        # 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.
@@ -332,6 +348,8 @@ 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),
                 )
                 nodes[element] = node
         return nodes

metaflow/lint.py

@@ -60,13 +60,41 @@ def check_reserved_words(graph):
 def check_basic_steps(graph):
     if graph.start_step is None:
         raise LintWarn(
-            "Your flow must have exactly one step with no incoming transitions "
-            "(the entry point). Found zero or multiple candidates."
+            "Your flow must have exactly one start step. Either name a step "
+            "'start' or use @step(start=True)."
         )
     if graph.end_step is None:
         raise LintWarn(
-            "Your flow must have exactly one step with no outgoing transitions "
-            "(the terminal step). Found zero or multiple candidates."
+            "Your flow must have exactly one end step. Either name a step "
+            "'end' or use @step(end=True)."
+        )
+
+
+@linter.ensure_static_graph
+@linter.check
+def check_start_end_degree(graph):
+    """Validate that the start step has no incoming and the end step has no outgoing."""
+    if graph.start_step is None or graph.end_step is None:
+        return
+
+    start_node = graph[graph.start_step]
+    if start_node.in_funcs:
+        raise LintWarn(
+            "The start step *%s* has incoming transitions from %s. "
+            "A start step must have no incoming transitions."
+            % (graph.start_step, ", ".join(start_node.in_funcs)),
+            start_node.func_lineno,
+            start_node.source_file,
+        )
+
+    end_node = graph[graph.end_step]
+    if end_node.out_funcs:
+        raise LintWarn(
+            "The end step *%s* has outgoing transitions. "
+            "An end step must have no outgoing transitions (no self.next())."
+            % graph.end_step,
+            end_node.func_lineno,
+            end_node.source_file,
         )
 
 

metaflow/parameters.py

@@ -520,13 +520,7 @@ def wrapper(cmd):
 
 
 class InitParameter(Parameter):
-    """
-    Parameter for StepSpec init() phase.
-
-    Behaves identically to Parameter in CLI mode. In direct invocation,
-    InitParameter values are passed to the constructor (alongside Config)
-    rather than to __call__.
-    """
+    """Parameter for StepSpec init() phase."""
 
     IS_INIT_PARAMETER = True