Prasanth345 commented on code in PR #53927:
URL: https://github.com/apache/airflow/pull/53927#discussion_r2250644038
##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -140,20 +140,42 @@ def execute_async(
del self.edge_queued_tasks[key]
self.validate_airflow_tasks_run_command(command) # type:
ignore[attr-defined]
- session.add(
- EdgeJobModel(
+
+ # Check if job already exists with same dag_id, task_id, run_id,
map_index, try_number
+ existing_job = (
+ session.query(EdgeJobModel)
+ .filter_by(
dag_id=key.dag_id,
task_id=key.task_id,
run_id=key.run_id,
map_index=key.map_index,
try_number=key.try_number,
- state=TaskInstanceState.QUEUED,
- queue=queue or DEFAULT_QUEUE,
- concurrency_slots=task_instance.pool_slots,
- command=str(command),
)
+ .first()
)
+ if existing_job:
+ # self.log.info(f"EdgeExecutor: Job already exists for {key},
updating it.")
+ existing_job.state = TaskInstanceState.QUEUED
+ existing_job.queue = queue or DEFAULT_QUEUE
+ existing_job.concurrency_slots = task_instance.pool_slots
+ existing_job.command = str(command)
+ else:
+ session.add(
+ EdgeJobModel(
+ dag_id=key.dag_id,
+ task_id=key.task_id,
+ run_id=key.run_id,
+ map_index=key.map_index,
+ try_number=key.try_number,
+ state=TaskInstanceState.QUEUED,
+ queue=queue or DEFAULT_QUEUE,
+ concurrency_slots=task_instance.pool_slots,
+ command=str(command),
+ )
+ )
+ # self.log.info(f"EdgeExecutor: Inserted new job for {key}")
Review Comment:
I’ve removed the commented-out debug logs to keep the code clean
##########
providers/edge3/src/airflow/providers/edge3/executors/edge_executor.py:
##########
@@ -168,20 +190,42 @@ def queue_workload(
task_instance = workload.ti
key = task_instance.key
- session.add(
- EdgeJobModel(
+
+ # Check if job already exists with same dag_id, task_id, run_id,
map_index, try_number
+ existing_job = (
+ session.query(EdgeJobModel)
+ .filter_by(
dag_id=key.dag_id,
task_id=key.task_id,
run_id=key.run_id,
map_index=key.map_index,
try_number=key.try_number,
- state=TaskInstanceState.QUEUED,
- queue=task_instance.queue,
- concurrency_slots=task_instance.pool_slots,
- command=workload.model_dump_json(),
)
+ .first()
)
+ if existing_job:
+ # self.log.info(f"EdgeExecutor: Workload already exists for {key},
updating it.")
+ existing_job.state = TaskInstanceState.QUEUED
+ existing_job.queue = task_instance.queue
+ existing_job.concurrency_slots = task_instance.pool_slots
+ existing_job.command = workload.model_dump_json()
+ else:
+ session.add(
+ EdgeJobModel(
+ dag_id=key.dag_id,
+ task_id=key.task_id,
+ run_id=key.run_id,
+ map_index=key.map_index,
+ try_number=key.try_number,
+ state=TaskInstanceState.QUEUED,
+ queue=task_instance.queue,
+ concurrency_slots=task_instance.pool_slots,
+ command=workload.model_dump_json(),
+ )
+ )
+ # self.log.info(f"EdgeExecutor: Inserted new workload for {key}")
Review Comment:
I’ve removed the commented-out debug logs to keep the code clean
--
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]
