Repository: PrefectHQ/prefect

Issue:
Move start_task_ids from FlowRunner to Environment
In order to facilitate passing `start_tasks` from external systems, we added a `start_task_ids` argument to `FlowRunner.run()`, which essentially dynamically creates the appropriate `start_tasks` argument. This means an external system (or command line) can use start_tasks without loading a Python Task and passing it to the FlowRunner.

This logic is really only used by `Environments`, which are the canonical way of interacting with flows from a remote system. I propose that it should therefore move to `Environment.run()` and leave `FlowRunner.run()` a cleaner API.

While this would be nominally a breaking change, it would have no breaking impact based on any current use of Prefect.


Discussion:

Resolving patch:
diff --git a/src/prefect/engine/flow_runner.py b/src/prefect/engine/flow_runner.py
--- a/src/prefect/engine/flow_runner.py
+++ b/src/prefect/engine/flow_runner.py
@@ -130,7 +130,6 @@ def run(
         state: State = None,
         task_states: Dict[Task, State] = None,
         start_tasks: Iterable[Task] = None,
-        start_task_ids: Iterable[str] = None,
         return_tasks: Iterable[Task] = None,
         parameters: Dict[str, Any] = None,
         task_runner_state_handlers: Iterable[Callable] = None,
@@ -149,8 +148,6 @@ def run(
             - start_tasks ([Task], optional): list of Tasks to begin computation
                 from; if any `start_tasks` have upstream dependencies, their states may
                 need to be provided as well.
-            - start_task_ids ([str], optional): equivalent to `start_tasks`, but accepts
-                a list of task IDs. The two options may be used simultaneously.
             - return_tasks ([Task], optional): list of Tasks to include in the
                 final returned Flow state. Defaults to `None`
             - parameters (dict, optional): dictionary of any needed Parameter
@@ -189,7 +186,6 @@ def run(
                     state,
                     task_states=task_states,
                     start_tasks=start_tasks,
-                    start_task_ids=start_task_ids,
                     return_tasks=return_tasks,
                     task_runner_state_handlers=task_runner_state_handlers,
                     executor=executor,
@@ -269,7 +265,6 @@ def get_flow_run_state(
         state: State,
         task_states: Dict[Task, State],
         start_tasks: Iterable[Task],
-        start_task_ids: Iterable[str],
         return_tasks: Set[Task],
         task_runner_state_handlers: Iterable[Callable],
         executor: "prefect.engine.executors.base.Executor",
@@ -285,8 +280,6 @@ def get_flow_run_state(
             - start_tasks ([Task], optional): list of Tasks to begin computation
                 from; if any `start_tasks` have upstream dependencies, their states may need to be provided as well.
                 Defaults to `self.flow.root_tasks()`
-            - start_task_ids ([str], optional): equivalent to `start_tasks`, but accepts
-                a list of task IDs. The two options may be used simultaneously.
             - return_tasks ([Task], optional): list of Tasks to include in the
                 final returned Flow state. Defaults to `None`
             - task_runner_state_handlers (Iterable[Callable], optional): A list of state change
@@ -307,9 +300,6 @@ def get_flow_run_state(
         # make copies to avoid modifying the user-supplied values
         task_states = dict(task_states or {})
         start_tasks = list(start_tasks or [])
-        if any(i not in self.flow.task_ids for i in start_task_ids or []):
-            raise ValueError("Invalid start_task_ids.")
-        start_tasks.extend(self.flow.task_ids[i] for i in start_task_ids or [])
 
         # don't make a copy to allow dynamic modification
         if return_tasks is None:
diff --git a/src/prefect/environments.py b/src/prefect/environments.py
--- a/src/prefect/environments.py
+++ b/src/prefect/environments.py
@@ -30,7 +30,7 @@
 import tempfile
 import textwrap
 import uuid
-from typing import Iterable
+from typing import Iterable, List
 
 import cloudpickle
 import docker
@@ -81,12 +81,16 @@ def build(self, flow: "prefect.Flow") -> "Environment":
         """
         raise NotImplementedError()
 
-    def run(self, runner_kwargs: dict = None) -> "prefect.engine.state.State":
+    def run(
+        self, start_task_ids: List[str] = None, runner_kwargs: dict = None
+    ) -> "prefect.engine.state.State":
         """
         Runs the `Flow` represented by this environment.
 
         Args:
-            - runner_kwargs (dict): Any arguments for `FlowRunner.run()`
+            - start_task_ids (List[str]): A list of Task ids that will be converted to the
+                `start_tasks` argument of the `FlowRunner`
+            - runner_kwargs (dict): Any arguments for `FlowRunner.run()`.
         """
         raise NotImplementedError()
 
@@ -182,22 +186,32 @@ def build(self, flow: "prefect.Flow") -> "LocalEnvironment":
             serialized_flow=self.serialize_flow_to_bytes(flow),
         )
 
-    def run(self, runner_kwargs: dict = None) -> "prefect.engine.state.State":
+    def run(
+        self, start_task_ids: List[str] = None, runner_kwargs: dict = None
+    ) -> "prefect.engine.state.State":
         """
         Runs the `Flow` represented by this environment.
 
         Args:
-            - runner_kwargs (dict): Any arguments for `FlowRunner.run()`
+            - start_task_ids (List[str]): A list of Task ids that will be converted to the
+                `start_tasks` argument of the `FlowRunner`
+            - runner_kwargs (dict): Any arguments for `FlowRunner.run()`.
         """
 
+        runner_kwargs = runner_kwargs or {}
+
         if not self.serialized_flow:
             raise ValueError(
                 "No serialized flow found! Has this environment been built?"
             )
         flow = self.deserialize_flow_from_bytes(self.serialized_flow)
+        if start_task_ids is not None:
+            if set(start_task_ids).difference(flow.task_ids):
+                raise ValueError("Invalid start_task_ids.")
+            runner_kwargs["start_tasks"] = [flow.task_ids[i] for i in start_task_ids]
         runner_cls = prefect.engine.get_default_flow_runner_class()
         runner = runner_cls(flow=flow)
-        return runner.run(**(runner_kwargs or {}))
+        return runner.run(**runner_kwargs)
 
 
 class ContainerEnvironment(Environment):
@@ -317,12 +331,16 @@ def build(
                 python_dependencies=self.python_dependencies,
             )
 
-    def run(self, runner_kwargs: dict = None) -> None:
+    def run(
+        self, start_task_ids: List[str] = None, runner_kwargs: dict = None
+    ) -> "prefect.engine.state.State":
         """
         Runs the `Flow` represented by this environment.
 
         Args:
-            - runner_kwargs (dict): Any arguments for `FlowRunner.run()`
+            - start_task_ids (List[str]): A list of Task ids that will be converted to the
+                `start_tasks` argument of the `FlowRunner`
+            - runner_kwargs (dict): Any arguments for `FlowRunner.run()`.
         """
 
         client = docker.APIClient(base_url="unix://var/run/docker.sock")


Test patch:
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -502,7 +502,6 @@ def test_running_state_finishes(self):
             state=Running(),
             task_states={},
             start_tasks=[],
-            start_task_ids=[],
             return_tasks=set(),
             task_runner_state_handlers=[],
             executor=LocalExecutor(),
@@ -519,7 +518,6 @@ def test_other_states_raise_endrun(self, state):
                 state=state,
                 task_states={},
                 start_tasks=[],
-                start_task_ids=[],
                 return_tasks=set(),
                 task_runner_state_handlers=[],
                 executor=Executor(),
@@ -535,7 +533,6 @@ def determine_final_state(self, *args, **kwargs):
             state=Running(),
             task_states={},
             start_tasks=[],
-            start_task_ids=[],
             return_tasks=set(),
             task_runner_state_handlers=[],
             executor=LocalExecutor(),
@@ -556,45 +553,6 @@ def test_start_tasks_are_respected(self):
         assert "Task not evaluated" in state.result[t2].message
         assert "Task not evaluated" not in state.result[t3].message
 
-    def test_start_tasks_ids_are_respected(self):
-        f = Flow()
-        t1, t2, t3 = Task(), Task(), Task()
-        f.add_edge(t1, t2)
-        f.add_edge(t2, t3)
-        state = FlowRunner(flow=f).run(
-            start_task_ids=[t3.id], return_tasks=[t1, t2, t3]
-        )
-
-        assert "Task not evaluated" in state.result[t1].message
-        assert "Task not evaluated" in state.result[t2].message
-        assert "Task not evaluated" not in state.result[t3].message
-
-    def test_start_tasks_ids_can_be_combined_with_start_tasks(self):
-        f = Flow()
-        t1, t2, t3 = Task(), Task(), Task()
-        f.add_edge(t1, t2)
-        f.add_edge(t2, t3)
-        state = FlowRunner(flow=f).run(
-            start_tasks=[t2], start_task_ids=[t3.id], return_tasks=[t1, t2, t3]
-        )
-
-        assert "Task not evaluated" in state.result[t1].message
-        assert "Task not evaluated" not in state.result[t2].message
-        assert "Task not evaluated" not in state.result[t3].message
-
-    def test_start_tasks_ids_can_overlap_with_start_tasks(self):
-        f = Flow()
-        t1, t2, t3 = Task(), Task(), Task()
-        f.add_edge(t1, t2)
-        f.add_edge(t2, t3)
-        state = FlowRunner(flow=f).run(
-            start_tasks=[t2], start_task_ids=[t2.id], return_tasks=[t1, t2, t3]
-        )
-
-        assert "Task not evaluated" in state.result[t1].message
-        assert "Task not evaluated" not in state.result[t2].message
-        assert "Task not evaluated" not in state.result[t3].message
-
     def test_invalid_start_task(self):
         f = Flow()
         t1 = Task()
@@ -605,16 +563,6 @@ def test_invalid_start_task(self):
         assert state.is_failed()
         assert "not found in Flow" in str(state.result)
 
-    def test_invalid_start_task_id(self):
-        f = Flow()
-        t1 = Task()
-        f.add_task(t1)
-
-        state = FlowRunner(flow=f).run(start_task_ids=["nope"])
-
-        assert state.is_failed()
-        assert "Invalid start_task_ids" in str(state.result)
-
     def test_start_tasks_ignores_triggers(self):
         f = Flow()
         t1, t2 = SuccessTask(), SuccessTask()
diff --git a/tests/test_environments.py b/tests/test_environments.py
--- a/tests/test_environments.py
+++ b/tests/test_environments.py
@@ -276,3 +276,26 @@ def test_serialize(self):
         env = LocalEnvironment().build(error_flow())
         s = env.serialize()
         assert isinstance(s, dict)
+
+    def test_start_task_ids(self):
+        @prefect.task
+        def error():
+            1 / 0
+
+        with prefect.Flow() as flow:
+            e = error()
+            x = prefect.Task()
+            x.set_upstream(e)
+
+        env = LocalEnvironment()
+        state = env.build(flow).run(start_task_ids=[x.id])
+        assert state.is_successful()
+
+    def test_invalid_start_task_ids(self):
+
+        flow = prefect.Flow()
+
+        env = LocalEnvironment()
+        with pytest.raises(ValueError) as exc:
+            env.build(flow).run(start_task_ids=["nope"])
+        assert "Invalid start_task_ids" in str(exc.value)