Update task engine to increment run_count when entering RUNNING s…

…tate (#15436)
PrefectHQ · Sep 19, 2024 · ab6c49d · ab6c49d
1 parent e248442
commit ab6c49d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/src/prefect/task_engine.py b/src/prefect/task_engine.py
@@ -411,6 +411,9 @@ def set_state(self, state: State, force: bool = False) -> State:
         self.task_run.state_type = new_state.type
         self.task_run.state_name = new_state.name
 
+        if new_state.is_running():
+            self.task_run.run_count += 1
+
         if new_state.is_final():
             if isinstance(state.data, BaseResult) and state.data.has_cached_object():
                 # Avoid fetching the result unless it is cached, otherwise we defeat
@@ -690,8 +693,9 @@ async def wait_until_ready(self):
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
             await anyio.sleep(sleep_time if sleep_time > 0 else 0)
+            new_state = Retrying() if self.state.name == "AwaitingRetry" else Running()
             self.set_state(
-                Retrying() if self.state.name == "AwaitingRetry" else Running(),
+                new_state,
                 force=True,
             )
 
@@ -773,7 +777,6 @@ def call_task_fn(
         if transaction.is_committed():
             result = transaction.read()
         else:
-            self.task_run.run_count += 1
             if self.task_run.tags:
                 # Acquire a concurrency slot for each tag, but only if a limit
                 # matching the tag already exists.
@@ -929,6 +932,9 @@ async def set_state(self, state: State, force: bool = False) -> State:
         self.task_run.state_type = new_state.type
         self.task_run.state_name = new_state.name
 
+        if new_state.is_running():
+            self.task_run.run_count += 1
+
         if new_state.is_final():
             if (
                 isinstance(new_state.data, BaseResult)
@@ -1196,8 +1202,9 @@ async def wait_until_ready(self):
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
             await anyio.sleep(sleep_time if sleep_time > 0 else 0)
+            new_state = Retrying() if self.state.name == "AwaitingRetry" else Running()
             await self.set_state(
-                Retrying() if self.state.name == "AwaitingRetry" else Running(),
+                new_state,
                 force=True,
             )
 
@@ -1280,7 +1287,6 @@ async def call_task_fn(
         if transaction.is_committed():
             result = transaction.read()
         else:
-            self.task_run.run_count += 1
             if self.task_run.tags:
                 # Acquire a concurrency slot for each tag, but only if a limit
                 # matching the tag already exists.

diff --git a/tests/events/client/instrumentation/test_task_run_state_change_events.py b/tests/events/client/instrumentation/test_task_run_state_change_events.py
@@ -139,7 +139,7 @@ def happy_path():
             },
             "flow_run_run_count": 1,
             "name": task_run.name,
-            "run_count": 0,
+            "run_count": 1,
             "tags": [],
             "task_inputs": {},
             "total_run_time": 0.0,
@@ -334,7 +334,7 @@ def happy_path():
             },
             "flow_run_run_count": 1,
             "name": task_run.name,
-            "run_count": 0,
+            "run_count": 1,
             "tags": [],
             "task_inputs": {},
             "total_run_time": 0.0,