shivaam commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3084566815
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -406,60 +470,60 @@ def process_queue(self, queue_url: str):
continue
return_code = body.get("return_code")
ser_task_key = body.get("task_key")
- # Fetch the real task key from the running_tasks dict, using the
serialized task key.
+ # Fetch the real task key from the running_workloads dict, using
the serialized task key.
try:
- task_key = self.running_tasks[ser_task_key]
+ workload_key = self.running_workloads[ser_task_key]
except KeyError:
self.log.debug(
- "Received task %s from the queue which is not found in
running tasks, it is likely "
+ "Received workload %s from the queue which is not found in
running workloads, it is likely "
"from another Lambda Executor sharing this queue or might
be a stale message that needs "
"deleting manually. Marking the message as visible again.",
ser_task_key,
)
- # Mark task as visible again in SQS so that another executor
can pick it up.
+ # Mark workload as visible again in SQS so that another
executor can pick it up.
self.sqs_client.change_message_visibility(
QueueUrl=queue_url,
ReceiptHandle=receipt_handle,
VisibilityTimeout=0,
)
continue
- if task_key:
+ if workload_key:
if return_code == 0:
- self.success(task_key)
+ self.success(cast("TaskInstanceKey", workload_key))
self.log.info(
- "Successful Lambda invocation for task %s received
from SQS queue.", task_key
+ "Successful Lambda invocation for workload %s received
from SQS queue.",
+ workload_key,
)
else:
- self.fail(task_key)
+ self.fail(cast("TaskInstanceKey", workload_key))
Review Comment:
@ferruzzi Just flagging — I'm doing the same cast in my ECS PR and I think
it only works by accident.
BaseExecutor.fail() is typed for TaskInstanceKey and writes
TaskInstanceState.FAILED into event_buffer. The scheduler routes callback keys
through a branch that checks against CallbackState.{RUNNING,FAILED,SUCCESS},
and that check only passes because both enums are str, Enum sharing the same
string values. So the cast silences mypy and runtime works, but it's relying on
two unrelated enums never drifting — and nothing enforces that alignment.
Possible approach we can take is to update BaseExecutor.fail() / success()
(and change_state) to branch on isinstance(key, TaskInstanceKey) and write the
correct enum per key type — TaskInstanceState for TI keys, CallbackState for
callback keys. The signature would widen to WorkloadKey, the cast disappears
from executor call sites, and the event buffer holds semantically correct
values.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
