just store task_func

Author: bmckerry

clients/python/src/taskbroker_client/worker/workerchild.py

@@ -59,13 +59,15 @@ class ActivationWithPendingFutures:
         execution_start_time: Timestamp of task execution start.
         futures_start_time: Timestamp of when the pending futures were enqueued.
         pending_futures: Set of pending futures generated by executing this activation.
+        task_func: The Task object related to this activation.
     """
 
     inflight: InflightTaskActivation
     status: TaskActivationStatus.ValueType
     execution_start_time: float
     futures_start_time: float
     pending_futures: set[ProducerFuture[BrokerValue[KafkaPayload]]]
+    task_func: Task[Any, Any]
 
 
 @contextlib.contextmanager
@@ -261,24 +263,23 @@ def await_task_futures(task: ActivationWithPendingFutures) -> None:
             and submitted for retry (if the policy allows).
             """
             RESULT_TIMEOUT_SEC = 1
-            task_func = _get_known_task(task.inflight.activation)
             try:
                 # We don't care about the actual result value,
                 # we just care if result() raises or not
                 [f.result(RESULT_TIMEOUT_SEC) for f in task.pending_futures]
             # If any pending producer futures failed, retry the task
             except Exception:
                 task.status = TASK_ACTIVATION_STATUS_FAILURE
-                if task_func and task_func.retry:
+                if task.task_func.retry:
                     retry_state = task.inflight.activation.retry_state
-                    if not task_func.retry.max_attempts_reached(retry_state):
+                    if not task.task_func.retry.max_attempts_reached(retry_state):
                         task.status = TASK_ACTIVATION_STATUS_RETRY
             pending_task_futures.remove(task)
             _task_execution_complete(
                 inflight=task.inflight,
                 next_state=task.status,
                 execution_start_time=task.execution_start_time,
-                task_func=task_func,
+                task_func=task.task_func,
                 futures_start_time=task.futures_start_time,
             )
 
@@ -460,6 +461,7 @@ def check_task_future_completion() -> None:
                     execution_start_time=execution_start_time,
                     futures_start_time=time.time(),
                     pending_futures=task_produced_futures,
+                    task_func=task_func,
                 )
                 pending_task_futures.append(pending_task)