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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
