[task scheduling] add intermediate Pending state and update policy (#…

…11985)

src/prefect/server/orchestration/core_policy.py

@@ -83,7 +83,7 @@ class AutonomousTaskPolicy(BaseOrchestrationPolicy):
 
     def priority():
         return [
-            PreventRunningToRunningTransitions,
+            PreventPendingTransitions,
             CacheRetrieval,
             HandleTaskTerminalStateTransitions,
             SecureTaskConcurrencySlots,  # retrieve cached states even if slots are full
@@ -1004,20 +1004,3 @@ async def before_transition(
                 state=None,
                 reason="This run has already made this state transition.",
             )
-
-
-class PreventRunningToRunningTransitions(BaseOrchestrationRule):
-    """Prevents transitions from Running to Running states."""
-
-    FROM_STATES = [StateType.RUNNING]
-    TO_STATES = [StateType.RUNNING]
-
-    async def before_transition(
-        self,
-        initial_state: Optional[states.State],
-        proposed_state: Optional[states.State],
-        context: OrchestrationContext,
-    ) -> None:
-        await self.abort_transition(
-            reason="Cannot transition from Running to Running.",
-        )

src/prefect/task_engine.py

@@ -9,6 +9,7 @@
 )
 
 import anyio
+import greenback
 from typing_extensions import Literal
 
 from prefect._internal.concurrency.api import create_call, from_async, from_sync
@@ -68,4 +69,5 @@ async def submit_autonomous_task_to_engine(
             if task.isasync:
                 return await from_async.wait_for_call_in_loop_thread(begin_run)
             else:
+                await greenback.ensure_portal()
                 return from_sync.wait_for_call_in_loop_thread(begin_run)

src/prefect/task_server.py

@@ -12,12 +12,14 @@
 from prefect._internal.concurrency.api import create_call, from_sync
 from prefect.client.schemas.objects import TaskRun
 from prefect.client.subscriptions import Subscription
+from prefect.engine import propose_state
 from prefect.logging.loggers import get_logger
 from prefect.results import ResultFactory
 from prefect.settings import (
     PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING,
     PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS,
 )
+from prefect.states import Pending
 from prefect.task_engine import submit_autonomous_task_to_engine
 from prefect.task_runners import BaseTaskRunner, ConcurrentTaskRunner
 from prefect.utilities.asyncutils import asyncnullcontext, sync_compatible
@@ -108,9 +110,9 @@ async def _subscribe_to_task_scheduling(self):
             [task.task_key for task in self.tasks],
         ):
             logger.info(f"Received task run: {task_run.id} - {task_run.name}")
-            await self._submit_pending_task_run(task_run)
+            await self._submit_scheduled_task_run(task_run)
 
-    async def _submit_pending_task_run(self, task_run: TaskRun):
+    async def _submit_scheduled_task_run(self, task_run: TaskRun):
         logger.debug(
             f"Found task run: {task_run.name!r} in state: {task_run.state.name!r}"
         )
@@ -152,6 +154,19 @@ async def _submit_pending_task_run(self, task_run: TaskRun):
             f"Submitting run {task_run.name!r} of task {task.name!r} to engine"
         )
 
+        state = await propose_state(
+            client=get_client(),  # TODO prove that we cannot use self._client here
+            state=Pending(),
+            task_run_id=task_run.id,
+        )
+
+        if not state.is_pending():
+            logger.warning(
+                f"Aborted task run {task_run.id!r} -"
+                f" server returned a non-pending state {state.type.value!r}."
+                " Task run may have already begun execution."
+            )
+
         self._runs_task_group.start_soon(
             partial(
                 submit_autonomous_task_to_engine,

tests/server/api/test_task_runs.py

@@ -526,8 +526,11 @@ async def test_set_task_run_state_returns_404_on_missing_flow_run(
         )
         assert response.status_code == status.HTTP_404_NOT_FOUND
 
-    async def test_autonomous_task_run_aborts_if_transitions_to_running_twice(
-        self, client, session
+    @pytest.mark.parametrize(
+        "incoming_state_type", ["PENDING", "RUNNING", "CANCELLED", "CANCELLING"]
+    )
+    async def test_autonomous_task_run_aborts_if_enters_pending_from_disallowed_state(
+        self, client, session, incoming_state_type
     ):
         autonomous_task_run = await models.task_runs.create_task_run(
             session=session,
@@ -543,7 +546,7 @@ async def test_autonomous_task_run_aborts_if_transitions_to_running_twice(
 
         response_1 = await client.post(
             f"/task_runs/{autonomous_task_run.id}/set_state",
-            json=dict(state=dict(type="RUNNING")),
+            json=dict(state=dict(type=incoming_state_type)),
         )
 
         api_response_1 = OrchestrationResult.parse_obj(response_1.json())
@@ -552,7 +555,7 @@ async def test_autonomous_task_run_aborts_if_transitions_to_running_twice(
 
         response_2 = await client.post(
             f"/task_runs/{autonomous_task_run.id}/set_state",
-            json=dict(state=dict(type="RUNNING")),
+            json=dict(state=dict(type="PENDING")),
         )
 
         api_response_2 = OrchestrationResult.parse_obj(response_2.json())

tests/server/orchestration/test_core_policy.py

@@ -31,7 +31,6 @@
     PreventDuplicateTransitions,
     PreventPendingTransitions,
     PreventRunningTasksFromStoppedFlows,
-    PreventRunningToRunningTransitions,
     ReleaseTaskConcurrencySlots,
     RenameReruns,
     RetryFailedFlows,
@@ -3134,20 +3133,3 @@ async def test_same_transition_id(
 
         # states have the same transition id so the transition should be rejected
         assert ctx.response_status == SetStateStatus.REJECT
-
-
-class TestPreventRunningToRunningTransitions:
-    async def test_prevents_running_to_running_transitions(
-        self,
-        session,
-        initialize_orchestration,
-    ):
-        transition = (StateType.RUNNING, StateType.RUNNING)
-        context = await initialize_orchestration(
-            session, "flow", *transition, initial_details=None, proposed_details=None
-        )
-
-        async with PreventRunningToRunningTransitions(context, *transition) as ctx:
-            await ctx.validate_proposed_state()
-
-        assert ctx.response_status == SetStateStatus.ABORT