jscheffl commented on code in PR #53927:
URL: https://github.com/apache/airflow/pull/53927#discussion_r2250097350
##########
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:
And not to forget, for performance you should write it with lazy string
formatting as `self.log.debug("EdgeExecutor: Inserted new job for %s", key)`
##########
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:
Can you remove the commended-out debug (or keep then and make them
log.debug() instead)?
##########
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:
Same here, log.debug() would be good.
--
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]
