You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "SameerMesiah97 (via GitHub)" <gi...@apache.org> on 2026/03/07 13:39:38 UTC
[PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
SameerMesiah97 opened a new pull request, #63035:
URL: https://github.com/apache/airflow/pull/63035
**Description**
This change adds support for callback workloads (`ExecuteCallback`) to the `AwsLambdaExecutor`.
Previously, the executor only supported task workloads (`ExecuteTask`). This update extends the executor to accept, queue, and execute callback workloads alongside task workloads. The executor now maintains a `queued_callbacks` structure and updates `queue_workload()` to register `ExecuteCallback` workloads using the callback identifier (`callback.id`) as the workload key.
The workload processing flow has been updated to support callbacks throughout the executor lifecycle. `_process_workloads()` now handles both `ExecuteTask` and `ExecuteCallback` workloads, dispatching them through the same execution pathway. `execute_async()` has been extended to serialize both workload types and forward them to the Lambda runtime using `python -m airflow.sdk.execution_time.execute_workload`. Additionally, `attempt_task_runs()` and related task tracking logic have been updated to support string-based workload identifiers for callbacks while maintaining JSON-serialized `TaskInstanceKey` identifiers for task workloads.
**Rationale**
The ExecutorCallback framework introduced in Airflow 3.2 allows executors to execute synchronous callbacks on worker infrastructure. Executors must therefore be able to accept and dispatch `ExecuteCallback` workloads in addition to task workloads.
`AwsLambdaExecutor` delegates execution to the Task SDK runtime by invoking `airflow.sdk.execution_time.execute_workload` inside the Lambda environment. As reflected in the direction of PR #62645, the expectation is that handling of both task and callback execution will ultimately occur within the Task SDK runtime rather than inside individual executors. Forwarding callback workloads to the same runtime entrypoint used for tasks aligns the Lambda executor with this model.
**Tests**
Added unit tests verifying that:
* Callback workloads are queued, processed, and dispatched to Lambda using the Task SDK runtime command.
* Callback workloads correctly propagate queue overrides specified in `callback.data["queue"]`.
* Callback workloads are correctly adopted using string-based `external_executor_id` values.
**Documentation**
Docstrings have been updated in `AwsLambdaExecutor` to describe support for both task and callback workloads.
**Backwards Compatibility**
This change preserves existing behavior for task workloads.
Task workloads continue to use JSON-serialized `TaskInstanceKey` identifiers for executor tracking and scheduler adoption. Callback workloads use their callback identifier string as the workload key.
Related: #62887
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2928032513
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4339282685
> Thanks for the continued engagement on this one @SameerMesiah97!
>
> Since @ferruzzi is out, I'm checking in on this one. I trust Dennis on the code review front, but have you done manual testing in a while? There have been a lot of changes. If you could do some quick UAT to ensure the code is al behaving how you expect then I'll be happy to merge! Also, testing with the oldest version of Airflow that providers support (2.11.0) would be helpful too, to test all the back compat logic.
I have used a production image for a local live deployment using the latest code in my local branch (rebased on top of main) to execute the following DAG successfully::
```
from datetime import datetime
from airflow import DAG
from airflow.operators.bash import BashOperator
def success_callback(context):
print("Callback executed! Task succeeded.")
print(f"Task: {context['task_instance_key_str']}")
with DAG(
dag_id="lambda_executor_smoke",
start_date=datetime(2024, 1, 1),
schedule=None,
catchup=False,
) as dag:
test_task = BashOperator(
task_id="test_task",
bash_command="echo 'Hello from Lambda Executor' && exit 0",
on_success_callback=success_callback
)
```
Additionally, I used the base image for 2.11.0 and installed a local distribution based on the amazon provider code in my local branch (rebased on top of main). I can confirm that there are no compatibility issues and that the backwards compatibility shims are functional.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114133728
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
+ self, workload: workloads.ExecuteTask | workloads.ExecuteCallback, session: Session
+ ) -> None:
+ try:
+ # Delegate to BaseExecutor (This will work in 3.3.0+).
+ super().queue_workload(workload, session)
+ except ValueError:
+ from airflow.executors import workloads
+
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
- if not isinstance(workload, workloads.ExecuteTask):
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
+ for workload in workload_items:
+ queue: str | None
+ key: WorkloadKey
+ command: CommandType
+ if isinstance(workload, workloads.ExecuteTask):
+ command = [workload]
+ key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
+
+ del self.queued_tasks[key]
+
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ command = [workload]
+ key = workload.callback.id
+ queue = None
+
+ if isinstance(workload.callback.data, dict) and "queue" in workload.callback.data:
+ queue = workload.callback.data["queue"]
+
+ del self.queued_callbacks[key]
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ self.running.add(key)
+ continue
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ def execute_async(
+ self,
+ key: WorkloadKey,
+ command: CommandType,
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
- :param command: The shell command string to execute.
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
+ :param command: The workload command or serialized shell command to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ from airflow.executors import workloads
+
+ workload = command[0]
+
+ if AIRFLOW_V_3_3_PLUS:
+ if not isinstance(workload, (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ ser_input = workload.model_dump_json()
+
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
- self.pending_tasks.append(
+ self.pending_workloads.append(
LambdaQueuedTask(
key, command, queue if queue else "", executor_config or {}, 1, timezone.utcnow()
)
)
- def attempt_task_runs(self):
+ def attempt_workload_runs(self):
Review Comment:
Okay. I left this intentionally because my understanding is that these are not public APIs. Just for my own curiosity, is there some reason why we are doing this that I am not aware of?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3118090328
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
> Why not key? It's there explicitly so there is parity with ti.key
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
> Why not key? It's there explicitly so there is parity with ti.key
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "shivaam (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114168263
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
You are right, that's my mistake sorry. We can't do that until min_ver hits 3.2. Let's go with one standard replacement across all of these though. Try this:
```python
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
if isinstance(workload, workloads.ExecuteTask):
self.queued_tasks[workload.ti.key] = workload
return
if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
self.queued_callbacks[workload.callback.key] = workload
return
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114051109
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
Pretty sure we can drop this method entirely and let BaseExecutor's version cover the work.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
+ self, workload: workloads.ExecuteTask | workloads.ExecuteCallback, session: Session
+ ) -> None:
+ try:
+ # Delegate to BaseExecutor (This will work in 3.3.0+).
+ super().queue_workload(workload, session)
+ except ValueError:
+ from airflow.executors import workloads
+
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
- if not isinstance(workload, workloads.ExecuteTask):
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
+ for workload in workload_items:
+ queue: str | None
+ key: WorkloadKey
+ command: CommandType
+ if isinstance(workload, workloads.ExecuteTask):
+ command = [workload]
+ key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
+
+ del self.queued_tasks[key]
+
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ command = [workload]
+ key = workload.callback.id
+ queue = None
+
+ if isinstance(workload.callback.data, dict) and "queue" in workload.callback.data:
+ queue = workload.callback.data["queue"]
+
+ del self.queued_callbacks[key]
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ self.running.add(key)
+ continue
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ def execute_async(
+ self,
+ key: WorkloadKey,
+ command: CommandType,
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
- :param command: The shell command string to execute.
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
+ :param command: The workload command or serialized shell command to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ from airflow.executors import workloads
+
+ workload = command[0]
+
+ if AIRFLOW_V_3_3_PLUS:
+ if not isinstance(workload, (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ ser_input = workload.model_dump_json()
+
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
- self.pending_tasks.append(
+ self.pending_workloads.append(
LambdaQueuedTask(
key, command, queue if queue else "", executor_config or {}, 1, timezone.utcnow()
)
)
- def attempt_task_runs(self):
+ def attempt_workload_runs(self):
Review Comment:
Here and elsewhere where you renamed the methods, I think you need a back-compat shim like this:
```python
def attempt_task_runs(self):
"""Deprecated: use attempt_workload_runs."""
warnings.warn(
"attempt_task_runs is deprecated, use attempt_workload_runs instead.",
DeprecationWarning,
stacklevel=2,
)
return self.attempt_workload_runs()
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3118090328
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
> Why not key? It's there explicitly so there is parity with ti.key
I tried using key but the test for callbacks started failing (do callbacks even have a ‘key’ attribute?). I will try again to see if I can get ‘key’ to work.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2927950734
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
+ command: CommandType | Sequence[workloads.All],
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
:param command: The shell command string to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ if AIRFLOW_V_3_2_PLUS:
+ if not isinstance(command[0], (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(command[0])}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(command[0], workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(command[0])}")
+
+ workload = command[0]
Review Comment:
`workload = command[0]` has been moved just below `if len(command) == 1 `otherwise it could throw an `IndexError` if `command` is an empty list. Not at the top of the method but ithe function should look cleaner now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4048711715
> > Thanks for the contribution! This looks like a good start 🙂 But it's the bare minimum changes. Please be sure to cleanup your code before submitting a PR. Make sure types are correct, all comments, variables, method names updated. It takes a long time for maintainers to review code so we want to be using that time wisely.
>
> @o-nikolas @ferruzzi
>
> This means the PR is not currently reviewable. I have converted it to draft until I can get into an acceptable state.
Wonderful, thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4329399694
Thanks for the continued engagement on this one @SameerMesiah97!
Since @ferruzzi is out, I'm checking in on this one. I trust Dennis on the code review front, but have you done manual testing in a while? There have been a lot of changes. If you could do some quick UAT to ensure the code is al behaving how you expect then I'll be happy to merge!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3127432871
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
> Why not key? It's there explicitly so there is parity with ti.key
key is being used for callbacks now. I have deleted my comment.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
@ferruzzi
I added the below. I believe it should be `workload.callback.id` instead of `workload.callback.key`. I am leaving this so that others are aware.
```
# TODO: Remove this once the minimum supported version is 3.2+, and defer to BaseExecutor.queue_workload.
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
if isinstance(workload, workloads.ExecuteTask):
self.queued_tasks[workload.ti.key] = workload
return
if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
self.queued_callbacks[workload.callback.id] = workload
return
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114307890
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
+ self, workload: workloads.ExecuteTask | workloads.ExecuteCallback, session: Session
+ ) -> None:
+ try:
+ # Delegate to BaseExecutor (This will work in 3.3.0+).
+ super().queue_workload(workload, session)
+ except ValueError:
+ from airflow.executors import workloads
+
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
- if not isinstance(workload, workloads.ExecuteTask):
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
+ for workload in workload_items:
+ queue: str | None
+ key: WorkloadKey
+ command: CommandType
+ if isinstance(workload, workloads.ExecuteTask):
+ command = [workload]
+ key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
+
+ del self.queued_tasks[key]
+
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ command = [workload]
+ key = workload.callback.id
+ queue = None
+
+ if isinstance(workload.callback.data, dict) and "queue" in workload.callback.data:
+ queue = workload.callback.data["queue"]
+
+ del self.queued_callbacks[key]
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ self.running.add(key)
+ continue
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ def execute_async(
+ self,
+ key: WorkloadKey,
+ command: CommandType,
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
- :param command: The shell command string to execute.
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
+ :param command: The workload command or serialized shell command to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ from airflow.executors import workloads
+
+ workload = command[0]
+
+ if AIRFLOW_V_3_3_PLUS:
+ if not isinstance(workload, (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ ser_input = workload.model_dump_json()
+
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
- self.pending_tasks.append(
+ self.pending_workloads.append(
LambdaQueuedTask(
key, command, queue if queue else "", executor_config or {}, 1, timezone.utcnow()
)
)
- def attempt_task_runs(self):
+ def attempt_workload_runs(self):
Review Comment:
> I'm perhaps being overly paranoid? It's possible someone is subclassing this for their own personal custom executors and it doesn't hurt to be defensive.
That's a great point. I agree.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4179335872
@SameerMesiah97 This PR has a few issues that need to be addressed before it can be reviewed — please see our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :warning: **Unresolved review comments**: This PR has 8 unresolved review threads from maintainers: **@o-nikolas** (MEMBER): 6 unresolved threads; **@ferruzzi** (MEMBER): 2 unresolved threads. Please review and resolve all inline review comments before requesting another review. You can resolve a conversation by clicking 'Resolve conversation' on each thread after addressing the feedback. See [pull request guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst).
> **Note:** Your branch is **360 commits behind `main`**. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.
**What to do next:**
- The comment informs you what you need to do.
- Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
- There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates.
- Maintainers will then proceed with a normal review.
There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3142359551
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -49,29 +51,39 @@
from airflow.executors import workloads
from airflow.models.taskinstance import TaskInstance
+ if AIRFLOW_V_3_3_PLUS:
+ from airflow.executors.workloads.types import WorkloadKey as _WorkloadKey
+
+ WorkloadKey: TypeAlias = _WorkloadKey
+ else:
+ WorkloadKey: TypeAlias = TaskInstanceKey # type: ignore[no-redef, misc]
+
class AwsLambdaExecutor(BaseExecutor):
"""
- An Airflow Executor that submits tasks to AWS Lambda asynchronously.
+ An Airflow Executor that submits workloads (tasks and callbacks) to AWS Lambda asynchronously.
When execute_async() is called, the executor invokes a specified AWS Lambda function (asynchronously)
- with a payload that includes the task command and a unique task key.
+ with a payload that includes the workload command and a unique workload key.
The Lambda function writes its result directly to an SQS queue, which is then polled by this executor
- to update task state in Airflow.
+ to update workload state in Airflow.
"""
supports_multi_team: bool = True
+ if AIRFLOW_V_3_3_PLUS:
+ supports_callbacks: bool = True
+
if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
# In the v3 path, we store workloads, not commands as strings.
- # TODO: TaskSDK: move this type change into BaseExecutor
- queued_tasks: dict[TaskInstanceKey, workloads.All] # type: ignore[assignment]
+ # TODO: TaskSDK: move this type change into BaseExecutor.
+ queued_tasks: dict[WorkloadKey, workloads.All] # type: ignore[assignment]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
- self.pending_tasks: deque = deque()
- self.running_tasks: dict[str, TaskInstanceKey] = {}
+ self.pending_workloads: deque = deque()
+ self.running_workloads: dict[str, WorkloadKey] = {}
Review Comment:
Backwards compatibility shims have been added..
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -49,29 +51,39 @@
from airflow.executors import workloads
from airflow.models.taskinstance import TaskInstance
+ if AIRFLOW_V_3_3_PLUS:
+ from airflow.executors.workloads.types import WorkloadKey as _WorkloadKey
+
+ WorkloadKey: TypeAlias = _WorkloadKey
+ else:
+ WorkloadKey: TypeAlias = TaskInstanceKey # type: ignore[no-redef, misc]
+
class AwsLambdaExecutor(BaseExecutor):
"""
- An Airflow Executor that submits tasks to AWS Lambda asynchronously.
+ An Airflow Executor that submits workloads (tasks and callbacks) to AWS Lambda asynchronously.
When execute_async() is called, the executor invokes a specified AWS Lambda function (asynchronously)
- with a payload that includes the task command and a unique task key.
+ with a payload that includes the workload command and a unique workload key.
The Lambda function writes its result directly to an SQS queue, which is then polled by this executor
- to update task state in Airflow.
+ to update workload state in Airflow.
"""
supports_multi_team: bool = True
+ if AIRFLOW_V_3_3_PLUS:
+ supports_callbacks: bool = True
+
if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
# In the v3 path, we store workloads, not commands as strings.
- # TODO: TaskSDK: move this type change into BaseExecutor
- queued_tasks: dict[TaskInstanceKey, workloads.All] # type: ignore[assignment]
+ # TODO: TaskSDK: move this type change into BaseExecutor.
+ queued_tasks: dict[WorkloadKey, workloads.All] # type: ignore[assignment]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
- self.pending_tasks: deque = deque()
- self.running_tasks: dict[str, TaskInstanceKey] = {}
+ self.pending_workloads: deque = deque()
+ self.running_workloads: dict[str, WorkloadKey] = {}
Review Comment:
Backwards compatibility shims have been added.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas merged PR #63035:
URL: https://github.com/apache/airflow/pull/63035
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2916522907
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/utils.py:
##########
@@ -66,5 +65,5 @@ class AllLambdaConfigKeys(InvokeLambdaKwargsConfigKeys):
END_WAIT_TIMEOUT = "end_wait_timeout"
-CommandType = Sequence[str]
+CommandType = Sequence[Any]
Review Comment:
I agree. I did it this way to pass MyPy checks. I will see if I can narrow it down, even if it requires me passing `# type: ignore`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2927967220
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -473,13 +531,12 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
]:
for ti, ser_task_key in serialized_task_keys:
try:
- task_key = TaskInstanceKey.from_dict(json.loads(ser_task_key))
+ data = json.loads(ser_task_key)
+ task_key = TaskInstanceKey.from_dict(data)
except Exception:
- # If that task fails to deserialize, we should just skip it.
- self.log.exception(
- "Task failed to be adopted because the key could not be deserialized"
- )
- continue
+ # Callback workloads use string keys.
+ task_key = ser_task_key
+
Review Comment:
I have kept the functional aspects of `try_adopt_task_instances` the same as in my first attempt for now. Once @ferruzzi confirms whether or not this needs to handle callbacks (even if it is just logging exceptions), I will adjust this accordingly.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4300642770
@ferruzzi
I have pushed my latest set of changes in response to your feedback.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2943273761
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/utils.py:
##########
@@ -66,5 +67,11 @@ class AllLambdaConfigKeys(InvokeLambdaKwargsConfigKeys):
END_WAIT_TIMEOUT = "end_wait_timeout"
-CommandType = Sequence[str]
+if AIRFLOW_V_3_2_PLUS:
+ from airflow.executors.workloads import ExecuteCallback, ExecuteTask
+
+ CommandType: TypeAlias = Sequence[str] | Sequence[ExecuteTask | ExecuteCallback]
Review Comment:
Thoughts for future work, feel free to read and resolve:
`Sequence[ExecuteTask | ExecuteCallback]`
If we see this pattern more, we may want to add another TypeAlias to workloads/tpes, maybe call it `ExecutorWorkload`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2936009236
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -473,13 +531,12 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
]:
for ti, ser_task_key in serialized_task_keys:
try:
- task_key = TaskInstanceKey.from_dict(json.loads(ser_task_key))
+ data = json.loads(ser_task_key)
+ task_key = TaskInstanceKey.from_dict(data)
except Exception:
- # If that task fails to deserialize, we should just skip it.
- self.log.exception(
- "Task failed to be adopted because the key could not be deserialized"
- )
- continue
+ # Callback workloads use string keys.
+ task_key = ser_task_key
+
Review Comment:
@ferruzzi So you are okay with keeping it as is? Or revert it to the state before the implementation with exception logging removed?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3127429571
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
+ self, workload: workloads.ExecuteTask | workloads.ExecuteCallback, session: Session
+ ) -> None:
+ try:
+ # Delegate to BaseExecutor (This will work in 3.3.0+).
+ super().queue_workload(workload, session)
+ except ValueError:
+ from airflow.executors import workloads
+
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
- if not isinstance(workload, workloads.ExecuteTask):
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
+ for workload in workload_items:
+ queue: str | None
+ key: WorkloadKey
+ command: CommandType
+ if isinstance(workload, workloads.ExecuteTask):
+ command = [workload]
+ key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
+
+ del self.queued_tasks[key]
+
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ command = [workload]
+ key = workload.callback.id
+ queue = None
+
+ if isinstance(workload.callback.data, dict) and "queue" in workload.callback.data:
+ queue = workload.callback.data["queue"]
+
+ del self.queued_callbacks[key]
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ self.running.add(key)
+ continue
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ def execute_async(
+ self,
+ key: WorkloadKey,
+ command: CommandType,
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
- :param command: The shell command string to execute.
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
+ :param command: The workload command or serialized shell command to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ from airflow.executors import workloads
+
+ workload = command[0]
+
+ if AIRFLOW_V_3_3_PLUS:
+ if not isinstance(workload, (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ ser_input = workload.model_dump_json()
+
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
- self.pending_tasks.append(
+ self.pending_workloads.append(
LambdaQueuedTask(
key, command, queue if queue else "", executor_config or {}, 1, timezone.utcnow()
)
)
- def attempt_task_runs(self):
+ def attempt_workload_runs(self):
Review Comment:
@ferruzzi
I had to use `AirflowProviderDeprecationWarning` to pass CI.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114149755
##########
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:
This should be resolved with a rebase now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4330496017
> Thanks for the continued engagement on this one @SameerMesiah97!
>
> Since @ferruzzi is out, I'm checking in on this one. I trust Dennis on the code review front, but have you done manual testing in a while? There have been a lot of changes. If you could do some quick UAT to ensure the code is al behaving how you expect then I'll be happy to merge! Also, testing with the oldest version of Airflow that providers support (2.11.0) would be helpful too, to test all the back compat logic.
I did my last manual test on April 19. I can of course do another but I believe the results from my previous should still be valid as the changes since then have been focused on backwards compatibility.
I will do another round of tests (especially with Airflow 2.11.0 like you requested) but it will take me a few days at the very least.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2955723250
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/utils.py:
##########
@@ -66,5 +67,11 @@ class AllLambdaConfigKeys(InvokeLambdaKwargsConfigKeys):
END_WAIT_TIMEOUT = "end_wait_timeout"
-CommandType = Sequence[str]
+if AIRFLOW_V_3_2_PLUS:
+ from airflow.executors.workloads import ExecuteCallback, ExecuteTask
+
+ CommandType: TypeAlias = Sequence[str] | Sequence[ExecuteTask | ExecuteCallback]
Review Comment:
I tried `Sequence[ExecuteTask | ExecuteCallback]` before (and just right now) but this was/is causing typing errors in `execute_async` as the argument for command is of type `Sequence[str]` in the legacy path.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2943205548
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,144 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
Review Comment:
Sorry I wasn't clear, the variable name was what I was referring to. I see that at least some of it was existing code, but we should clean that up, not just fix your new changes.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2943254916
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -313,59 +377,59 @@ def attempt_task_runs(self):
failure_reasons.append(str(e))
if failure_reasons:
- # Make sure the number of attempts does not exceed max invoke attempts
+ # Make sure the number of attempts does not exceed max invoke attempts.
if int(attempt_number) < int(self.max_invoke_attempts):
- task_to_run.attempt_number += 1
- task_to_run.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
+ workload_to_run.attempt_number += 1
+ workload_to_run.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
attempt_number
)
- self.pending_tasks.append(task_to_run)
+ self.pending_workloads.append(workload_to_run)
else:
reasons_str = ", ".join(failure_reasons)
self.log.error(
"Lambda invoke %s has failed a maximum of %s times. Marking as failed. Reasons: %s",
- task_key,
+ workload_key,
attempt_number,
reasons_str,
)
self.log_task_event(
event="lambda invoke failure",
- ti_key=task_key,
+ ti_key=workload_key,
Review Comment:
Yup, good catch. I think the cleanest solution for now will be to add an optional `workload` parameter to Log() which will work for either type, and a note that we need to clean that up in the future. Don't remove the task_instance parameter since that will break things, but add the new `workload` and use it here.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114127330
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
> Pretty sure we can drop this method entirely and let BaseExecutor's version cover the work.
That is what I initially tried but then I saw that it caused the backwards compatiblity tests for pre-3.20 to fail. For example, `queue_workload` in 3.0.6 looks like this:
```
def queue_workload(self, workload: workloads.All, session: Session) -> None:
raise ValueError(f"Un-handled workload kind {type(workload).__name__!r} in {type(self).__name__}")
```
Technically, the `AwsLambdaExecutor` is part of the Amazon provider package so I opted for the above approach to ensure backwards compatiblity (the CI failures led me there as well). I have no issues with letting `BaseExecutor` handle `queue_workload` but I am not sure how it can be done without breaking the backwards compatibility tests.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2936001006
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -49,29 +49,39 @@
from airflow.executors import workloads
from airflow.models.taskinstance import TaskInstance
+ if AIRFLOW_V_3_2_PLUS:
+ from airflow.executors.workloads.types import WorkloadKey as _WorkloadKey
+
+ WorkloadKey: TypeAlias = _WorkloadKey
+ else:
+ WorkloadKey: TypeAlias = TaskInstanceKey # type: ignore[no-redef, misc]
Review Comment:
`error: Cannot assign multiple types to name "WorkloadKey" without an explicit "Type[...]" annotation [misc]`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2927995217
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -411,7 +467,7 @@ def process_queue(self, queue_url: str):
task_key = self.running_tasks[ser_task_key]
except KeyError:
Review Comment:
Where appropriate I have changed 'task' to 'workload' in variable names, method signatures, comments, docstrings and log messsages. However, some variables I thought it best to leave them as is because it would be misleading to change them. For example, `ser_task_key` is derived from the value whose key name is `task_key` in the JSON payload, which constitutes the SQS message. It would not make sense to change `ser_task_key` to `ser_workload_key` just for the sake of consistency.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2928018969
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -411,7 +467,7 @@ def process_queue(self, queue_url: str):
task_key = self.running_tasks[ser_task_key]
except KeyError:
Review Comment:
That's interesting yeah. I suppose that key should really be `workload_key` in the return schema for the lambda executor. It should still continue to accept task, but we should try to migrate towards using workload there. But I agree that is out of scope for this PR :slightly_smiling_face:
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2927956642
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -461,6 +517,8 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
"""
Adopt task instances which have an external_executor_id (the serialized task key).
+ The external_executor_id may contain either a serialized TaskInstanceKey or a callback identifier string.
Review Comment:
The docstring has been made more clear. It now explains that `WorkloadKey` will be used from 3.2 onwards.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2955778709
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/utils.py:
##########
@@ -66,5 +67,11 @@ class AllLambdaConfigKeys(InvokeLambdaKwargsConfigKeys):
END_WAIT_TIMEOUT = "end_wait_timeout"
-CommandType = Sequence[str]
+if AIRFLOW_V_3_2_PLUS:
+ from airflow.executors.workloads import ExecuteCallback, ExecuteTask
+
+ CommandType: TypeAlias = Sequence[str] | Sequence[ExecuteTask | ExecuteCallback]
Review Comment:
Sorry, I just meant that we may want a new TypeAlias for that particular combo if we end up using it more. so this would end up being `CommandType: TypeAlias = Sequence[str | ExecutorWorkload]` or something like that, that's all.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2956215522
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,144 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
Review Comment:
> Sorry I wasn't clear, the variable name was what I was referring to. I see that at least some of it was existing code, but we should clean that up, not just fix your new changes.
Fixed in latest commit.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4354508260
> > Thanks for the continued engagement on this one @SameerMesiah97!
> > Since @ferruzzi is out, I'm checking in on this one. I trust Dennis on the code review front, but have you done manual testing in a while? There have been a lot of changes. If you could do some quick UAT to ensure the code is al behaving how you expect then I'll be happy to merge! Also, testing with the oldest version of Airflow that providers support (2.11.0) would be helpful too, to test all the back compat logic.
>
> @o-nikolas
>
> I have used a production image for a local live deployment using the latest code in my local branch (rebased on top of main) to execute the following DAG successfully::
>
> ```
> from datetime import datetime
>
> from airflow import DAG
> from airflow.operators.bash import BashOperator
>
>
> def success_callback(context):
> print("Callback executed! Task succeeded.")
> print(f"Task: {context['task_instance_key_str']}")
>
> with DAG(
> dag_id="lambda_executor_smoke",
> start_date=datetime(2024, 1, 1),
> schedule=None,
> catchup=False,
> ) as dag:
>
> test_task = BashOperator(
> task_id="test_task",
> bash_command="echo 'Hello from Lambda Executor' && exit 0",
> on_success_callback=success_callback
> )
> ```
>
> Additionally, I used the base image for 2.11.0 and installed a local distribution based on the amazon provider code in my local branch (rebased on top of main). I can confirm that there are no compatibility issues and that the backwards compatibility shims are functional.
Lovely! Thanks for doing the testing :smiley:
Given @ferruzzi's "approval"
>I am heading on vacation tomorrow so you can consider this my approval once that fix is done..
I'll approve and merge!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4017328159
@ferruzzi
Tagging you here as requested for tracking.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2921362632
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
Review Comment:
`WorkloadKey`
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -473,13 +531,12 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
]:
for ti, ser_task_key in serialized_task_keys:
try:
- task_key = TaskInstanceKey.from_dict(json.loads(ser_task_key))
+ data = json.loads(ser_task_key)
+ task_key = TaskInstanceKey.from_dict(data)
except Exception:
- # If that task fails to deserialize, we should just skip it.
- self.log.exception(
- "Task failed to be adopted because the key could not be deserialized"
- )
- continue
+ # Callback workloads use string keys.
+ task_key = ser_task_key
+
Review Comment:
@ferruzzi can callbacks be adopted? If not, then we don't need these changes here, or if anything, only to discard callbacks (we probably don't want to log for each one).
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -263,7 +313,7 @@ def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None,
def attempt_task_runs(self):
Review Comment:
Rename this method?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -339,20 +395,20 @@ def attempt_task_runs(self):
self.fail(task_key)
else:
status_code = response.get("StatusCode")
- self.log.info("Invoked Lambda for task %s with status %s", task_key, status_code)
+ self.log.info("Invoked Lambda for workload %s with status %s", task_key, status_code)
self.running_tasks[ser_task_key] = task_key
# Add the serialized task key as the info, this will be assigned on the ti as the external_executor_id
self.running_state(task_key, ser_task_key)
def sync_running_tasks(self):
Review Comment:
Should this method be renamed?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -411,7 +467,7 @@ def process_queue(self, queue_url: str):
task_key = self.running_tasks[ser_task_key]
except KeyError:
Review Comment:
Shouldn't all these mentions of `task(s)` in variables and comments be updated?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -426,12 +482,12 @@ def process_queue(self, queue_url: str):
if task_key:
if return_code == 0:
- self.success(task_key)
+ self.success(task_key) # type: ignore[arg-type]
Review Comment:
Why are we just ignoring this instead of getting the typing right?
Here and just below as well.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2914876060
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/utils.py:
##########
@@ -66,5 +65,5 @@ class AllLambdaConfigKeys(InvokeLambdaKwargsConfigKeys):
END_WAIT_TIMEOUT = "end_wait_timeout"
-CommandType = Sequence[str]
+CommandType = Sequence[Any]
Review Comment:
Can we narrow this down any?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/utils.py:
##########
@@ -37,9 +36,9 @@
@dataclass
class LambdaQueuedTask:
- """Represents a Lambda task that is queued. The task will be run in the next heartbeat."""
+ """Represents a Lambda workload that is queued. The task will be run in the next heartbeat."""
- key: TaskInstanceKey
+ key: TaskInstanceKey | str
Review Comment:
Please use `WorkloadKey` here from `airflow-core/src/airflow/executors/workloads/types.py`
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -461,6 +517,8 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
"""
Adopt task instances which have an external_executor_id (the serialized task key).
+ The external_executor_id may contain either a serialized TaskInstanceKey or a callback identifier string.
Review Comment:
WorkloadKey?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -71,7 +76,10 @@ class AwsLambdaExecutor(BaseExecutor):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.pending_tasks: deque = deque()
- self.running_tasks: dict[str, TaskInstanceKey] = {}
+ self.running_tasks: dict[str, TaskInstanceKey | str] = {}
Review Comment:
Same as below, `TaskInstanceKey | str` should be `executors.workloads.types.WorkloadKey` instead
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -179,9 +187,9 @@ def load_connections(self, check_connection: bool = True):
def sync(self):
"""
- Sync the executor with the current state of tasks.
+ Sync the executor with the current state of workloads.
- Check in on currently running tasks and attempt to run any new tasks that have been queued.
+ Check in on currently running workloads and attempt to run any new workloads that have been queued.
Review Comment:
Here and elsewhere: As Niko pointed out in one of my recent PRs, the user may not be familiar with the term `workload` yet so it would be beneficial to expand it a little so we can teach them. Consider:
```suggestion
Check in on currently running tasks and callbacks and attempt to run any new workloads that have been queued.
```
That way we're kind of teaching them what a `workload` is without making them go look it up.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
+ command: CommandType | Sequence[workloads.All],
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
:param command: The shell command string to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ if AIRFLOW_V_3_2_PLUS:
+ if not isinstance(command[0], (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(command[0])}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(command[0], workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(command[0])}")
+
+ workload = command[0]
Review Comment:
Consider moving this `workload = command[0]` to the top of the method and using `workload` throughout the method, it would clean things up a bit.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
+ command: CommandType | Sequence[workloads.All],
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
:param command: The shell command string to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ if AIRFLOW_V_3_2_PLUS:
+ if not isinstance(command[0], (workloads.ExecuteTask, workloads.ExecuteCallback)):
Review Comment:
can this be `if not isinstance(command[0], SchedulerWorkload)`, using the SchedulerWorkload from executors/workloads/types.py? If not, then leave this for now, @anishgirianish has a really cool PR in he works to simplify this.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114484953
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
@ferruzzi
I added the below. I believe it should be `workload.callback.id` instead of `workload.callback.key`. I am leaving this so that others are aware.
```
# TODO: Remove this once the minimum supported version is 3.2+, and defer to BaseExecutor.queue_workload.
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
if isinstance(workload, workloads.ExecuteTask):
self.queued_tasks[workload.ti.key] = workload
return
if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
self.queued_callbacks[workload.callback.id] = workload
return
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2927935033
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
+ command: CommandType | Sequence[workloads.All],
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
:param command: The shell command string to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ if AIRFLOW_V_3_2_PLUS:
+ if not isinstance(command[0], (workloads.ExecuteTask, workloads.ExecuteCallback)):
Review Comment:
Unfortunately, `SchedulerWorkload` is a `TypeAlias` so `isinstance` will not accept that. It has been left as is for now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2928027127
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -339,20 +395,20 @@ def attempt_task_runs(self):
self.fail(task_key)
else:
status_code = response.get("StatusCode")
- self.log.info("Invoked Lambda for task %s with status %s", task_key, status_code)
+ self.log.info("Invoked Lambda for workload %s with status %s", task_key, status_code)
self.running_tasks[ser_task_key] = task_key
# Add the serialized task key as the info, this will be assigned on the ti as the external_executor_id
self.running_state(task_key, ser_task_key)
def sync_running_tasks(self):
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114133728
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
+ self, workload: workloads.ExecuteTask | workloads.ExecuteCallback, session: Session
+ ) -> None:
+ try:
+ # Delegate to BaseExecutor (This will work in 3.3.0+).
+ super().queue_workload(workload, session)
+ except ValueError:
+ from airflow.executors import workloads
+
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
- if not isinstance(workload, workloads.ExecuteTask):
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
+ for workload in workload_items:
+ queue: str | None
+ key: WorkloadKey
+ command: CommandType
+ if isinstance(workload, workloads.ExecuteTask):
+ command = [workload]
+ key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
+
+ del self.queued_tasks[key]
+
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ command = [workload]
+ key = workload.callback.id
+ queue = None
+
+ if isinstance(workload.callback.data, dict) and "queue" in workload.callback.data:
+ queue = workload.callback.data["queue"]
+
+ del self.queued_callbacks[key]
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ self.running.add(key)
+ continue
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ def execute_async(
+ self,
+ key: WorkloadKey,
+ command: CommandType,
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
- :param command: The shell command string to execute.
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
+ :param command: The workload command or serialized shell command to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ from airflow.executors import workloads
+
+ workload = command[0]
+
+ if AIRFLOW_V_3_3_PLUS:
+ if not isinstance(workload, (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ ser_input = workload.model_dump_json()
+
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
- self.pending_tasks.append(
+ self.pending_workloads.append(
LambdaQueuedTask(
key, command, queue if queue else "", executor_config or {}, 1, timezone.utcnow()
)
)
- def attempt_task_runs(self):
+ def attempt_workload_runs(self):
Review Comment:
Okay. I left this intentionally because my understanding is that these are not public APIs. I will fix it.
But just for my own curiosity, is there some reason why we are doing this that I am not aware of? Is this for backwards compatibility with airflow core?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3118090328
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
> Why not key? It's there explicitly so there is parity with ti.key
I tried using key but the test for callbacks started failing which I fixed by using id instead. I will try again to see if I can get ‘key’ to work.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114138328
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
+ self, workload: workloads.ExecuteTask | workloads.ExecuteCallback, session: Session
+ ) -> None:
+ try:
+ # Delegate to BaseExecutor (This will work in 3.3.0+).
+ super().queue_workload(workload, session)
+ except ValueError:
+ from airflow.executors import workloads
+
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
- if not isinstance(workload, workloads.ExecuteTask):
raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
+ for workload in workload_items:
+ queue: str | None
+ key: WorkloadKey
+ command: CommandType
+ if isinstance(workload, workloads.ExecuteTask):
+ command = [workload]
+ key = workload.ti.key
+ queue = workload.ti.queue
+ executor_config = workload.ti.executor_config or {}
+
+ del self.queued_tasks[key]
+
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
+
+ if AIRFLOW_V_3_3_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ command = [workload]
+ key = workload.callback.id
+ queue = None
+
+ if isinstance(workload.callback.data, dict) and "queue" in workload.callback.data:
+ queue = workload.callback.data["queue"]
+
+ del self.queued_callbacks[key]
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ self.running.add(key)
+ continue
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ def execute_async(
+ self,
+ key: WorkloadKey,
+ command: CommandType,
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
- :param command: The shell command string to execute.
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
+ :param command: The workload command or serialized shell command to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ from airflow.executors import workloads
+
+ workload = command[0]
+
+ if AIRFLOW_V_3_3_PLUS:
+ if not isinstance(workload, (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ ser_input = workload.model_dump_json()
+
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
- self.pending_tasks.append(
+ self.pending_workloads.append(
LambdaQueuedTask(
key, command, queue if queue else "", executor_config or {}, 1, timezone.utcnow()
)
)
- def attempt_task_runs(self):
+ def attempt_workload_runs(self):
Review Comment:
I'm perhaps being overly paranoid? It's possible someone is subclassing this for their own personal custom executors and it doesn't hurt to be defensive.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2932890691
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -473,13 +531,12 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
]:
for ti, ser_task_key in serialized_task_keys:
try:
- task_key = TaskInstanceKey.from_dict(json.loads(ser_task_key))
+ data = json.loads(ser_task_key)
+ task_key = TaskInstanceKey.from_dict(data)
except Exception:
- # If that task fails to deserialize, we should just skip it.
- self.log.exception(
- "Task failed to be adopted because the key could not be deserialized"
- )
- continue
+ # Callback workloads use string keys.
+ task_key = ser_task_key
+
Review Comment:
Sorry for the delay. I don't believe they support adoption yet, but that might/should/could be added in the future (not this PR, obviously).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2933182418
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -411,7 +467,7 @@ def process_queue(self, queue_url: str):
task_key = self.running_tasks[ser_task_key]
except KeyError:
Review Comment:
Agreed on all points.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
+ command: CommandType | Sequence[workloads.All],
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
:param command: The shell command string to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ if AIRFLOW_V_3_2_PLUS:
+ if not isinstance(command[0], (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(command[0])}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(command[0], workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(command[0])}")
+
+ workload = command[0]
Review Comment:
Yup, that's great. I didn't mean literally first line, just up here where it can be resued throughout. 👍
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,144 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
Review Comment:
I agree with Niko. I know this was existing code, but that was a miss in an earlier review, `w` is a terrible variable name. Can you fix it while you are in here please?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -49,29 +49,39 @@
from airflow.executors import workloads
from airflow.models.taskinstance import TaskInstance
+ if AIRFLOW_V_3_2_PLUS:
+ from airflow.executors.workloads.types import WorkloadKey as _WorkloadKey
+
+ WorkloadKey: TypeAlias = _WorkloadKey
+ else:
+ WorkloadKey: TypeAlias = TaskInstanceKey # type: ignore[no-redef, misc]
Review Comment:
For my own learning, what's the `misc` error we're ignoring here? no-redef makes sense.
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,144 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
+ for w in workload_items:
+ queue: str | None
+ key: WorkloadKey
+ command: CommandType
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
+
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ del self.queued_callbacks[key]
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ self.running.add(key)
+ continue
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: WorkloadKey,
+ command: CommandType,
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
- :param command: The shell command string to execute.
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
+ :param command: The workload command or serialized shell command to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ from airflow.executors import workloads
+
+ workload = command[0]
+
+ if AIRFLOW_V_3_2_PLUS:
+ if not isinstance(workload, (workloads.ExecuteTask, workloads.ExecuteCallback)):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
else:
- raise RuntimeError(
- f"LambdaExecutor doesn't know how to handle workload of type: {type(command[0])}"
- )
+ if not isinstance(workload, workloads.ExecuteTask):
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- self.pending_tasks.append(
+ ser_input = workload.model_dump_json()
+
+ command = [
+ "python",
+ "-m",
+ "airflow.sdk.execution_time.execute_workload",
+ "--json-string",
+ ser_input,
+ ]
+
+ self.pending_workloads.append(
LambdaQueuedTask(
key, command, queue if queue else "", executor_config or {}, 1, timezone.utcnow()
)
)
- def attempt_task_runs(self):
+ def attempt_workload_runs(self):
"""
- Attempt to run tasks that are queued in the pending_tasks.
+ Attempt to run workloads that are queued in the pending_workloads.
- Each task is submitted to AWS Lambda with a payload containing the task key and command.
- The task key is used to track the task's state in Airflow.
+ Each workload is submitted to AWS Lambda with a payload containing the workload key and command.
+ The workload key is used to track the workload's state in Airflow.
"""
- queue_len = len(self.pending_tasks)
+ queue_len = len(self.pending_workloads)
for _ in range(queue_len):
- task_to_run = self.pending_tasks.popleft()
- task_key = task_to_run.key
- cmd = task_to_run.command
- attempt_number = task_to_run.attempt_number
+ workload_to_run = self.pending_workloads.popleft()
+ workload_key = workload_to_run.key
+ cmd = workload_to_run.command
+ attempt_number = workload_to_run.attempt_number
failure_reasons = []
- ser_task_key = json.dumps(task_key._asdict())
+
+ try:
+ ser_workload_key = json.dumps(workload_key._asdict())
+ except AttributeError:
+ # Callback workloads use string id.
+ ser_workload_key = workload_key
+
payload = {
- "task_key": ser_task_key,
+ "task_key": ser_workload_key,
Review Comment:
hrm. I guess eventually we'll want to update that payload key too, but this may not be the time for it... that feels like a bigger breaking change
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -313,59 +377,59 @@ def attempt_task_runs(self):
failure_reasons.append(str(e))
if failure_reasons:
- # Make sure the number of attempts does not exceed max invoke attempts
+ # Make sure the number of attempts does not exceed max invoke attempts.
if int(attempt_number) < int(self.max_invoke_attempts):
- task_to_run.attempt_number += 1
- task_to_run.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
+ workload_to_run.attempt_number += 1
+ workload_to_run.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
attempt_number
)
- self.pending_tasks.append(task_to_run)
+ self.pending_workloads.append(workload_to_run)
else:
reasons_str = ", ".join(failure_reasons)
self.log.error(
"Lambda invoke %s has failed a maximum of %s times. Marking as failed. Reasons: %s",
- task_key,
+ workload_key,
attempt_number,
reasons_str,
)
self.log_task_event(
event="lambda invoke failure",
- ti_key=task_key,
+ ti_key=workload_key,
Review Comment:
Non-blocking: Looks like we missed this parameter name in base executor... we'll have to get that later.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3083251289
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -468,19 +532,18 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
with Stats.timer("lambda_executor.adopt_task_instances.duration"):
adopted_tis: list[TaskInstance] = []
- if serialized_task_keys := [
+ if serialized_workload_keys := [
(ti, ti.external_executor_id) for ti in tis if ti.external_executor_id
]:
- for ti, ser_task_key in serialized_task_keys:
+ for ti, ser_workload_key in serialized_workload_keys:
try:
- task_key = TaskInstanceKey.from_dict(json.loads(ser_task_key))
+ data = json.loads(ser_workload_key)
+ workload_key = TaskInstanceKey.from_dict(data)
except Exception:
- # If that task fails to deserialize, we should just skip it.
- self.log.exception(
- "Task failed to be adopted because the key could not be deserialized"
- )
- continue
- self.running_tasks[ser_task_key] = task_key
+ # Callback workloads use string keys.
+ workload_key = ser_workload_key
Review Comment:
Can you catch the specific exception that you receive when the string keys fail to loads above? Then return the broad Exception case to logging an error message?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2936009236
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -473,13 +531,12 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
]:
for ti, ser_task_key in serialized_task_keys:
try:
- task_key = TaskInstanceKey.from_dict(json.loads(ser_task_key))
+ data = json.loads(ser_task_key)
+ task_key = TaskInstanceKey.from_dict(data)
except Exception:
- # If that task fails to deserialize, we should just skip it.
- self.log.exception(
- "Task failed to be adopted because the key could not be deserialized"
- )
- continue
+ # Callback workloads use string keys.
+ task_key = ser_task_key
+
Review Comment:
@ferruzzi So you are okay with keeping it as is? Or revert it to the state before the implementation i.e. log exception on failure?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4042939985
> Thanks for the contribution! This looks like a good start 🙂 But it's the bare minimum changes. Please be sure to cleanup your code before submitting a PR. Make sure types are correct, all comments, variables, method names updated. It takes a long time for maintainers to review code so we want to be using that time wisely.
@o-nikolas @ferruzzi
This means the PR is not currently reviewable. I have converted it to draft until I can get into an acceptable state.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4051506181
@ferruzzi @o-nikolas
I have just pushed my latest changes. Here is a brief summary of the last commit:
1. All variable names, method signatures, docstrings and comments have been updated where appropriate. For some sections of the code where there is some ambiguity (such as `try_adopt_task_instances`), I have left the naming as is. Please refer to my comments responding to your feedback.
2. The interaction between strict mypy conventions and the multiple branches needed to support different Airflow versions made explicit typing difficult in a few places. Where possible, I avoided `# type: ignore`, but in some cases using it was the clearest option.
3. You may notice the diff size increased as most of the tests had to be updated to accomodating the changes in variable names and method signatures. Also, while searching for occurrences of "task(s)" to replace with `workload(s)`, I also cleaned up surrounding comments (typos, capitalization, redundant wording, missing punctuation), which further increased the diff size.
CI is green now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2943211522
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -49,29 +49,39 @@
from airflow.executors import workloads
from airflow.models.taskinstance import TaskInstance
+ if AIRFLOW_V_3_2_PLUS:
+ from airflow.executors.workloads.types import WorkloadKey as _WorkloadKey
+
+ WorkloadKey: TypeAlias = _WorkloadKey
+ else:
+ WorkloadKey: TypeAlias = TaskInstanceKey # type: ignore[no-redef, misc]
Review Comment:
Huh., Thanks. I would have thought no-redef would have covered that, but I guess not. Thanks for humouring me.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2927924871
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -179,9 +187,9 @@ def load_connections(self, check_connection: bool = True):
def sync(self):
"""
- Sync the executor with the current state of tasks.
+ Sync the executor with the current state of workloads.
- Check in on currently running tasks and attempt to run any new tasks that have been queued.
+ Check in on currently running workloads and attempt to run any new workloads that have been queued.
Review Comment:
Done.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2928028066
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -263,7 +313,7 @@ def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None,
def attempt_task_runs(self):
Review Comment:
Done.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2928025684
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -426,12 +482,12 @@ def process_queue(self, queue_url: str):
if task_key:
if return_code == 0:
- self.success(task_key)
+ self.success(task_key) # type: ignore[arg-type]
Review Comment:
I have removed the `type :ignore` in favour of using `cast`. Now, arguably, it's another way of doing the same thing but `self.success` and `self.fail` are defined in `BaseExecutor` (which I am not going to touch in this PR) so this limits what can be done to satisfy typing whilst keeping the code explicit. If you have an alternative suggestion in mind, I am all ears.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2932597893
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -205,55 +213,97 @@ def sync(self):
self.log.exception("An error occurred while syncing tasks")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+
+ for w in workload_items:
+ key: TaskInstanceKey | str
+ command: list[workloads.All]
+ queue: str | None
+ if isinstance(w, workloads.ExecuteTask):
+ command = [w]
+ key = w.ti.key
+ queue = w.ti.queue
+ executor_config = w.ti.executor_config or {}
+
+ del self.queued_tasks[key]
- def _process_workloads(self, workloads: Sequence[workloads.All]) -> None:
- from airflow.executors.workloads import ExecuteTask
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ executor_config=executor_config,
+ )
+
+ self.running.add(key)
+ continue
- for w in workloads:
- if not isinstance(w, ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+ if AIRFLOW_V_3_2_PLUS and isinstance(w, workloads.ExecuteCallback):
+ command = [w]
+ key = w.callback.id
+ queue = None
- command = [w]
- key = w.ti.key
- queue = w.ti.queue
- executor_config = w.ti.executor_config or {}
+ if isinstance(w.callback.data, dict) and "queue" in w.callback.data:
+ queue = w.callback.data["queue"]
- del self.queued_tasks[key]
- self.execute_async(key=key, command=command, queue=queue, executor_config=executor_config) # type: ignore[arg-type]
- self.running.add(key)
+ del self.queued_callbacks[key]
- def execute_async(self, key: TaskInstanceKey, command: CommandType, queue=None, executor_config=None):
+ self.execute_async(
+ key=key,
+ command=command,
+ queue=queue,
+ )
+
+ self.running.add(key)
+ continue
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(w)}")
+
+ def execute_async(
+ self,
+ key: TaskInstanceKey | str,
+ command: CommandType | Sequence[workloads.All],
+ queue=None,
+ executor_config=None,
+ ):
"""
- Save the task to be executed in the next sync by inserting the commands into a queue.
+ Save the workload to be executed in the next sync by inserting the commands into a queue.
- :param key: A unique task key (typically a tuple identifying the task instance).
+ :param key: Unique workload key. Task workloads use TaskInstanceKey, callback workloads use a string id.
:param command: The shell command string to execute.
:param executor_config: (Unused) to keep the same signature as the base.
:param queue: (Unused) to keep the same signature as the base.
"""
if len(command) == 1:
- from airflow.executors.workloads import ExecuteTask
-
- if isinstance(command[0], ExecuteTask):
- workload = command[0]
- ser_input = workload.model_dump_json()
- command = [
- "python",
- "-m",
- "airflow.sdk.execution_time.execute_workload",
- "--json-string",
- ser_input,
- ]
+ if AIRFLOW_V_3_2_PLUS:
+ if not isinstance(command[0], (workloads.ExecuteTask, workloads.ExecuteCallback)):
Review Comment:
You are absolutely right, sorry about that. We could use BaseDagBundleWorkload... but let's leave it for now, I guess.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2936018521
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,144 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
Review Comment:
@ferruzzi
Can you be more specific about what you are agreeing on with Niko? I am not sure what action is required from me on thos comment besides changing the variable name?
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,144 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
from airflow.executors import workloads
- if not isinstance(workload, workloads.ExecuteTask):
- raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
- ti = workload.ti
- self.queued_tasks[ti.key] = workload
+ if isinstance(workload, workloads.ExecuteTask):
+ ti = workload.ti
+ self.queued_tasks[ti.key] = workload
+ return
+
+ if AIRFLOW_V_3_2_PLUS and isinstance(workload, workloads.ExecuteCallback):
+ self.queued_callbacks[workload.callback.id] = workload
+ return
+
+ raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}")
+
+ def _process_workloads(self, workload_items: Sequence[workloads.All]) -> None:
+ from airflow.executors import workloads
+
Review Comment:
@ferruzzi
Can you be more specific about what you are agreeing on with Niko? I am not sure what action is required from me on this comment besides changing the variable name?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "shivaam (via GitHub)" <gi...@apache.org>.
shivaam commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r2937496249
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -313,59 +377,59 @@ def attempt_task_runs(self):
failure_reasons.append(str(e))
if failure_reasons:
- # Make sure the number of attempts does not exceed max invoke attempts
+ # Make sure the number of attempts does not exceed max invoke attempts.
if int(attempt_number) < int(self.max_invoke_attempts):
- task_to_run.attempt_number += 1
- task_to_run.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
+ workload_to_run.attempt_number += 1
+ workload_to_run.next_attempt_time = timezone.utcnow() + calculate_next_attempt_delay(
attempt_number
)
- self.pending_tasks.append(task_to_run)
+ self.pending_workloads.append(workload_to_run)
else:
reasons_str = ", ".join(failure_reasons)
self.log.error(
"Lambda invoke %s has failed a maximum of %s times. Marking as failed. Reasons: %s",
- task_key,
+ workload_key,
attempt_number,
reasons_str,
)
self.log_task_event(
event="lambda invoke failure",
- ti_key=task_key,
+ ti_key=workload_key,
Review Comment:
When a callback workload exceeds max submit attempts, log_task_event is called with ti_key=workload_key. For callbacks, this key is a string UUID, not a TaskInstanceKey named tuple, which will cause errors since Log(task_instance=...) expects a TaskInstanceKey. Is this a problem @ferruzzi ?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3086325210
##########
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:
Yes. To give you more context, I only did this because I did not want to touch `BaseExecutor` in this PR but it most certainly should be modified to handle callbacks keys within `self.fail ()` and `self.success ()` in a follow up PR.
##########
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:
Yes. To give you more context, I only did this because I did not want to touch `BaseExecutor` in this PR but it most certainly should be modified to handle callback keys within `self.fail ()` and `self.success ()` in a follow up PR.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #63035:
URL: https://github.com/apache/airflow/pull/63035#issuecomment-4276695432
> Also @SameerMesiah97 Have you done any manual QA? Actually running the executor with some tasks and callbacks with these changes?
I have just tested the following DAG with a live cloud deployment using the `AwsLambdaExecutor`:
```
from datetime import datetime
from airflow import DAG
from airflow.operators.bash import BashOperator
def success_callback(context):
print("Callback executed! Task succeeded.")
print(f"Task: {context['task_instance_key_str']}")
with DAG(
dag_id="lambda_executor_smoke",
start_date=datetime(2024, 1, 1),
schedule=None,
catchup=False,
) as dag:
test_task = BashOperator(
task_id="test_task",
bash_command="echo 'Hello from Lambda Executor' && exit 0",
on_success_callback=success_callback
)
```
I can verify that both the task and callback executed successfully via the UI and the Cloudwatch logs.
> This looks close now! Just one more comment from me. Then I'm okay with the changes. But @ferruzzi should have the final say since he is more aware of the changes needed to support Callbacks.
Note that the above test was carried out before a new set of changes I will push now to address your feedback and accomodate changes in PR #65392.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3114663138
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -202,90 +212,150 @@ def sync(self):
"AWS credentials are either missing or expired: %s.\nRetrying connection", error
)
except Exception:
- self.log.exception("An error occurred while syncing tasks")
+ self.log.exception("An error occurred while syncing workloads.")
- def queue_workload(self, workload: workloads.All, session: Session | None) -> None:
- from airflow.executors import workloads
+ def queue_workload(
Review Comment:
Why not key? It's there explicitly so there is parity with ti.key
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add ExecuteCallback support to AwsLambdaExecutor [airflow]
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #63035:
URL: https://github.com/apache/airflow/pull/63035#discussion_r3133852800
##########
providers/amazon/src/airflow/providers/amazon/aws/executors/aws_lambda/lambda_executor.py:
##########
@@ -49,29 +51,39 @@
from airflow.executors import workloads
from airflow.models.taskinstance import TaskInstance
+ if AIRFLOW_V_3_3_PLUS:
+ from airflow.executors.workloads.types import WorkloadKey as _WorkloadKey
+
+ WorkloadKey: TypeAlias = _WorkloadKey
+ else:
+ WorkloadKey: TypeAlias = TaskInstanceKey # type: ignore[no-redef, misc]
+
class AwsLambdaExecutor(BaseExecutor):
"""
- An Airflow Executor that submits tasks to AWS Lambda asynchronously.
+ An Airflow Executor that submits workloads (tasks and callbacks) to AWS Lambda asynchronously.
When execute_async() is called, the executor invokes a specified AWS Lambda function (asynchronously)
- with a payload that includes the task command and a unique task key.
+ with a payload that includes the workload command and a unique workload key.
The Lambda function writes its result directly to an SQS queue, which is then polled by this executor
- to update task state in Airflow.
+ to update workload state in Airflow.
"""
supports_multi_team: bool = True
+ if AIRFLOW_V_3_3_PLUS:
+ supports_callbacks: bool = True
+
if TYPE_CHECKING and AIRFLOW_V_3_0_PLUS:
# In the v3 path, we store workloads, not commands as strings.
- # TODO: TaskSDK: move this type change into BaseExecutor
- queued_tasks: dict[TaskInstanceKey, workloads.All] # type: ignore[assignment]
+ # TODO: TaskSDK: move this type change into BaseExecutor.
+ queued_tasks: dict[WorkloadKey, workloads.All] # type: ignore[assignment]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
- self.pending_tasks: deque = deque()
- self.running_tasks: dict[str, TaskInstanceKey] = {}
+ self.pending_workloads: deque = deque()
+ self.running_workloads: dict[str, WorkloadKey] = {}
Review Comment:
Can you add property shim for `pending_tasks` and `running_tasks` as well? They'll look something like
```python
@property
def pending_tasks(self) -> deque:
"""Deprecated: use pending_workloads."""
return self.pending_workloads
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org