29 lines
1.3 KiB
Python
29 lines
1.3 KiB
Python
# Version: 2025-08-20-02
|
|
# This file contains custom hooks for the Airflow environment.
|
|
from airflow import settings
|
|
from airflow.models.dagrun import DagRun
|
|
from airflow.utils.session import provide_session
|
|
|
|
|
|
@provide_session
|
|
def task_instance_mutation_hook(ti, session=None):
|
|
if ti.dag_id == 'ytdlp_ops_worker_per_url':
|
|
# Query the DagRun from the DB using run_id to reliably get the conf.
|
|
# The ti.dag_run attribute is not always populated when the hook is called.
|
|
dag_run = session.query(DagRun).filter(DagRun.run_id == ti.run_id).first()
|
|
conf = dag_run.conf if dag_run else {}
|
|
worker_queue = conf.get('worker_queue')
|
|
|
|
if worker_queue:
|
|
print(f"MUTATION HOOK: For dag '{ti.dag_id}', pinning task '{ti.task_id}' (run_id: {ti.run_id}) to queue '{worker_queue}'.")
|
|
ti.queue = worker_queue
|
|
else:
|
|
print(f"MUTATION HOOK: For dag '{ti.dag_id}', no 'worker_queue' in conf for run_id '{ti.run_id}'. Falling back to 'queue-dl'.")
|
|
ti.queue = 'queue-dl'
|
|
|
|
# Register the hook only in appropriate contexts
|
|
# This hook can cause issues with the Triggerer, which does not have a `dag_run` context
|
|
# when it runs its own maintenance tasks.
|
|
if not settings.CONFIG.get('core', 'executor').lower().startswith('debug'):
|
|
settings.task_instance_mutation_hook = task_instance_mutation_hook
|