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