-
Notifications
You must be signed in to change notification settings - Fork 13.7k
/
dagrun.py
1246 lines (1111 loc) · 49.5 KB
/
dagrun.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import itertools
import os
import warnings
from collections import defaultdict
from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
Generator,
Iterable,
List,
NamedTuple,
Optional,
Sequence,
Set,
Tuple,
Union,
cast,
)
from sqlalchemy import (
Boolean,
Column,
ForeignKey,
Index,
Integer,
PickleType,
String,
UniqueConstraint,
and_,
func,
or_,
text,
)
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy.orm import joinedload, relationship, synonym
from sqlalchemy.orm.session import Session
from sqlalchemy.sql.expression import false, select, true
from airflow import settings
from airflow.callbacks.callback_requests import DagCallbackRequest
from airflow.configuration import conf as airflow_conf
from airflow.exceptions import AirflowException, TaskNotFound
from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
from airflow.models.mappedoperator import MappedOperator
from airflow.models.taskinstance import TaskInstance as TI
from airflow.models.tasklog import LogTemplate
from airflow.stats import Stats
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.dependencies_states import SCHEDULEABLE_STATES
from airflow.utils import timezone
from airflow.utils.helpers import is_container
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.sqlalchemy import UtcDateTime, nulls_first, skip_locked, tuple_in_condition, with_row_locks
from airflow.utils.state import DagRunState, State, TaskInstanceState
from airflow.utils.types import NOTSET, ArgNotSet, DagRunType
if TYPE_CHECKING:
from airflow.models.dag import DAG
from airflow.models.operator import Operator
class TISchedulingDecision(NamedTuple):
"""Type of return for DagRun.task_instance_scheduling_decisions"""
tis: List[TI]
schedulable_tis: List[TI]
changed_tis: bool
unfinished_tis: List[TI]
finished_tis: List[TI]
class DagRun(Base, LoggingMixin):
"""
DagRun describes an instance of a Dag. It can be created
by the scheduler (for regular runs) or by an external trigger
"""
__tablename__ = "dag_run"
id = Column(Integer, primary_key=True)
dag_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
queued_at = Column(UtcDateTime)
execution_date = Column(UtcDateTime, default=timezone.utcnow, nullable=False)
start_date = Column(UtcDateTime)
end_date = Column(UtcDateTime)
_state = Column('state', String(50), default=State.QUEUED)
run_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False)
creating_job_id = Column(Integer)
external_trigger = Column(Boolean, default=True)
run_type = Column(String(50), nullable=False)
conf = Column(PickleType)
# These two must be either both NULL or both datetime.
data_interval_start = Column(UtcDateTime)
data_interval_end = Column(UtcDateTime)
# When a scheduler last attempted to schedule TIs for this DagRun
last_scheduling_decision = Column(UtcDateTime)
dag_hash = Column(String(32))
# Foreign key to LogTemplate. DagRun rows created prior to this column's
# existence have this set to NULL. Later rows automatically populate this on
# insert to point to the latest LogTemplate entry.
log_template_id = Column(
Integer,
ForeignKey("log_template.id", name="task_instance_log_template_id_fkey", ondelete="NO ACTION"),
default=select([func.max(LogTemplate.__table__.c.id)]),
)
# Remove this `if` after upgrading Sphinx-AutoAPI
if not TYPE_CHECKING and "BUILDING_AIRFLOW_DOCS" in os.environ:
dag: "Optional[DAG]"
else:
dag: "Optional[DAG]" = None
__table_args__ = (
Index('dag_id_state', dag_id, _state),
UniqueConstraint('dag_id', 'execution_date', name='dag_run_dag_id_execution_date_key'),
UniqueConstraint('dag_id', 'run_id', name='dag_run_dag_id_run_id_key'),
Index('idx_last_scheduling_decision', last_scheduling_decision),
Index('idx_dag_run_dag_id', dag_id),
Index(
'idx_dag_run_running_dags',
'state',
'dag_id',
postgresql_where=text("state='running'"),
mssql_where=text("state='running'"),
sqlite_where=text("state='running'"),
),
# since mysql lacks filtered/partial indices, this creates a
# duplicate index on mysql. Not the end of the world
Index(
'idx_dag_run_queued_dags',
'state',
'dag_id',
postgresql_where=text("state='queued'"),
mssql_where=text("state='queued'"),
sqlite_where=text("state='queued'"),
),
)
task_instances = relationship(
TI, back_populates="dag_run", cascade='save-update, merge, delete, delete-orphan'
)
DEFAULT_DAGRUNS_TO_EXAMINE = airflow_conf.getint(
'scheduler',
'max_dagruns_per_loop_to_schedule',
fallback=20,
)
def __init__(
self,
dag_id: Optional[str] = None,
run_id: Optional[str] = None,
queued_at: Union[datetime, None, ArgNotSet] = NOTSET,
execution_date: Optional[datetime] = None,
start_date: Optional[datetime] = None,
external_trigger: Optional[bool] = None,
conf: Optional[Any] = None,
state: Optional[DagRunState] = None,
run_type: Optional[str] = None,
dag_hash: Optional[str] = None,
creating_job_id: Optional[int] = None,
data_interval: Optional[Tuple[datetime, datetime]] = None,
):
if data_interval is None:
# Legacy: Only happen for runs created prior to Airflow 2.2.
self.data_interval_start = self.data_interval_end = None
else:
self.data_interval_start, self.data_interval_end = data_interval
self.dag_id = dag_id
self.run_id = run_id
self.execution_date = execution_date
self.start_date = start_date
self.external_trigger = external_trigger
self.conf = conf or {}
if state is not None:
self.state = state
if queued_at is NOTSET:
self.queued_at = timezone.utcnow() if state == State.QUEUED else None
else:
self.queued_at = queued_at
self.run_type = run_type
self.dag_hash = dag_hash
self.creating_job_id = creating_job_id
super().__init__()
def __repr__(self):
return (
'<DagRun {dag_id} @ {execution_date}: {run_id}, state:{state}, '
'queued_at: {queued_at}. externally triggered: {external_trigger}>'
).format(
dag_id=self.dag_id,
execution_date=self.execution_date,
run_id=self.run_id,
state=self.state,
queued_at=self.queued_at,
external_trigger=self.external_trigger,
)
@property
def logical_date(self) -> datetime:
return self.execution_date
def get_state(self):
return self._state
def set_state(self, state: DagRunState):
if state not in State.dag_states:
raise ValueError(f"invalid DagRun state: {state}")
if self._state != state:
self._state = state
self.end_date = timezone.utcnow() if self._state in State.finished else None
if state == State.QUEUED:
self.queued_at = timezone.utcnow()
@declared_attr
def state(self):
return synonym('_state', descriptor=property(self.get_state, self.set_state))
@provide_session
def refresh_from_db(self, session: Session = NEW_SESSION) -> None:
"""
Reloads the current dagrun from the database
:param session: database session
"""
dr = session.query(DagRun).filter(DagRun.dag_id == self.dag_id, DagRun.run_id == self.run_id).one()
self.id = dr.id
self.state = dr.state
@classmethod
@provide_session
def active_runs_of_dags(cls, dag_ids=None, only_running=False, session=None) -> Dict[str, int]:
"""Get the number of active dag runs for each dag."""
query = session.query(cls.dag_id, func.count('*'))
if dag_ids is not None:
# 'set' called to avoid duplicate dag_ids, but converted back to 'list'
# because SQLAlchemy doesn't accept a set here.
query = query.filter(cls.dag_id.in_(list(set(dag_ids))))
if only_running:
query = query.filter(cls.state == State.RUNNING)
else:
query = query.filter(cls.state.in_([State.RUNNING, State.QUEUED]))
query = query.group_by(cls.dag_id)
return {dag_id: count for dag_id, count in query.all()}
@classmethod
def next_dagruns_to_examine(
cls,
state: DagRunState,
session: Session,
max_number: Optional[int] = None,
):
"""
Return the next DagRuns that the scheduler should attempt to schedule.
This will return zero or more DagRun rows that are row-level-locked with a "SELECT ... FOR UPDATE"
query, you should ensure that any scheduling decisions are made in a single transaction -- as soon as
the transaction is committed it will be unlocked.
:rtype: list[airflow.models.DagRun]
"""
from airflow.models.dag import DagModel
if max_number is None:
max_number = cls.DEFAULT_DAGRUNS_TO_EXAMINE
# TODO: Bake this query, it is run _A lot_
query = (
session.query(cls)
.filter(cls.state == state, cls.run_type != DagRunType.BACKFILL_JOB)
.join(DagModel, DagModel.dag_id == cls.dag_id)
.filter(DagModel.is_paused == false(), DagModel.is_active == true())
)
if state == State.QUEUED:
# For dag runs in the queued state, we check if they have reached the max_active_runs limit
# and if so we drop them
running_drs = (
session.query(DagRun.dag_id, func.count(DagRun.state).label('num_running'))
.filter(DagRun.state == DagRunState.RUNNING)
.group_by(DagRun.dag_id)
.subquery()
)
query = query.outerjoin(running_drs, running_drs.c.dag_id == DagRun.dag_id).filter(
func.coalesce(running_drs.c.num_running, 0) < DagModel.max_active_runs
)
query = query.order_by(
nulls_first(cls.last_scheduling_decision, session=session),
cls.execution_date,
)
if not settings.ALLOW_FUTURE_EXEC_DATES:
query = query.filter(DagRun.execution_date <= func.now())
return with_row_locks(
query.limit(max_number), of=cls, session=session, **skip_locked(session=session)
)
@classmethod
@provide_session
def find(
cls,
dag_id: Optional[Union[str, List[str]]] = None,
run_id: Optional[Iterable[str]] = None,
execution_date: Optional[Union[datetime, Iterable[datetime]]] = None,
state: Optional[DagRunState] = None,
external_trigger: Optional[bool] = None,
no_backfills: bool = False,
run_type: Optional[DagRunType] = None,
session: Session = NEW_SESSION,
execution_start_date: Optional[datetime] = None,
execution_end_date: Optional[datetime] = None,
) -> List["DagRun"]:
"""
Returns a set of dag runs for the given search criteria.
:param dag_id: the dag_id or list of dag_id to find dag runs for
:param run_id: defines the run id for this dag run
:param run_type: type of DagRun
:param execution_date: the execution date
:param state: the state of the dag run
:param external_trigger: whether this dag run is externally triggered
:param no_backfills: return no backfills (True), return all (False).
Defaults to False
:param session: database session
:param execution_start_date: dag run that was executed from this date
:param execution_end_date: dag run that was executed until this date
"""
qry = session.query(cls)
dag_ids = [dag_id] if isinstance(dag_id, str) else dag_id
if dag_ids:
qry = qry.filter(cls.dag_id.in_(dag_ids))
if is_container(run_id):
qry = qry.filter(cls.run_id.in_(run_id))
elif run_id is not None:
qry = qry.filter(cls.run_id == run_id)
if is_container(execution_date):
qry = qry.filter(cls.execution_date.in_(execution_date))
elif execution_date is not None:
qry = qry.filter(cls.execution_date == execution_date)
if execution_start_date and execution_end_date:
qry = qry.filter(cls.execution_date.between(execution_start_date, execution_end_date))
elif execution_start_date:
qry = qry.filter(cls.execution_date >= execution_start_date)
elif execution_end_date:
qry = qry.filter(cls.execution_date <= execution_end_date)
if state:
qry = qry.filter(cls.state == state)
if external_trigger is not None:
qry = qry.filter(cls.external_trigger == external_trigger)
if run_type:
qry = qry.filter(cls.run_type == run_type)
if no_backfills:
qry = qry.filter(cls.run_type != DagRunType.BACKFILL_JOB)
return qry.order_by(cls.execution_date).all()
@classmethod
@provide_session
def find_duplicate(
cls,
dag_id: str,
run_id: str,
execution_date: datetime,
session: Session = NEW_SESSION,
) -> Optional['DagRun']:
"""
Return an existing run for the DAG with a specific run_id or execution_date.
*None* is returned if no such DAG run is found.
:param dag_id: the dag_id to find duplicates for
:param run_id: defines the run id for this dag run
:param execution_date: the execution date
:param session: database session
"""
return (
session.query(cls)
.filter(
cls.dag_id == dag_id,
or_(cls.run_id == run_id, cls.execution_date == execution_date),
)
.one_or_none()
)
@staticmethod
def generate_run_id(run_type: DagRunType, execution_date: datetime) -> str:
"""Generate Run ID based on Run Type and Execution Date"""
return f"{run_type}__{execution_date.isoformat()}"
@provide_session
def get_task_instances(
self,
state: Optional[Iterable[Optional[TaskInstanceState]]] = None,
session: Session = NEW_SESSION,
) -> List[TI]:
"""Returns the task instances for this dag run"""
tis = (
session.query(TI)
.options(joinedload(TI.dag_run))
.filter(
TI.dag_id == self.dag_id,
TI.run_id == self.run_id,
)
)
if state:
if isinstance(state, str):
tis = tis.filter(TI.state == state)
else:
# this is required to deal with NULL values
if State.NONE in state:
if all(x is None for x in state):
tis = tis.filter(TI.state.is_(None))
else:
not_none_state = [s for s in state if s]
tis = tis.filter(or_(TI.state.in_(not_none_state), TI.state.is_(None)))
else:
tis = tis.filter(TI.state.in_(state))
if self.dag and self.dag.partial:
tis = tis.filter(TI.task_id.in_(self.dag.task_ids))
return tis.all()
@provide_session
def get_task_instance(
self,
task_id: str,
session: Session = NEW_SESSION,
*,
map_index: int = -1,
) -> Optional[TI]:
"""
Returns the task instance specified by task_id for this dag run
:param task_id: the task id
:param session: Sqlalchemy ORM Session
"""
return (
session.query(TI)
.filter_by(dag_id=self.dag_id, run_id=self.run_id, task_id=task_id, map_index=map_index)
.one_or_none()
)
def get_dag(self) -> "DAG":
"""
Returns the Dag associated with this DagRun.
:return: DAG
"""
if not self.dag:
raise AirflowException(f"The DAG (.dag) for {self} needs to be set")
return self.dag
@provide_session
def get_previous_dagrun(
self, state: Optional[DagRunState] = None, session: Session = NEW_SESSION
) -> Optional['DagRun']:
"""The previous DagRun, if there is one"""
filters = [
DagRun.dag_id == self.dag_id,
DagRun.execution_date < self.execution_date,
]
if state is not None:
filters.append(DagRun.state == state)
return session.query(DagRun).filter(*filters).order_by(DagRun.execution_date.desc()).first()
@provide_session
def get_previous_scheduled_dagrun(self, session: Session = NEW_SESSION) -> Optional['DagRun']:
"""The previous, SCHEDULED DagRun, if there is one"""
return (
session.query(DagRun)
.filter(
DagRun.dag_id == self.dag_id,
DagRun.execution_date < self.execution_date,
DagRun.run_type != DagRunType.MANUAL,
)
.order_by(DagRun.execution_date.desc())
.first()
)
@provide_session
def update_state(
self, session: Session = NEW_SESSION, execute_callbacks: bool = True
) -> Tuple[List[TI], Optional[DagCallbackRequest]]:
"""
Determines the overall state of the DagRun based on the state
of its TaskInstances.
:param session: Sqlalchemy ORM Session
:param execute_callbacks: Should dag callbacks (success/failure, SLA etc) be invoked
directly (default: true) or recorded as a pending request in the ``callback`` property
:return: Tuple containing tis that can be scheduled in the current loop & `callback` that
needs to be executed
"""
# Callback to execute in case of Task Failures
callback: Optional[DagCallbackRequest] = None
start_dttm = timezone.utcnow()
self.last_scheduling_decision = start_dttm
with Stats.timer(f"dagrun.dependency-check.{self.dag_id}"):
dag = self.get_dag()
info = self.task_instance_scheduling_decisions(session)
tis = info.tis
schedulable_tis = info.schedulable_tis
changed_tis = info.changed_tis
finished_tis = info.finished_tis
unfinished_tis = info.unfinished_tis
none_depends_on_past = all(not t.task.depends_on_past for t in unfinished_tis)
none_task_concurrency = all(t.task.max_active_tis_per_dag is None for t in unfinished_tis)
none_deferred = all(t.state != State.DEFERRED for t in unfinished_tis)
if unfinished_tis and none_depends_on_past and none_task_concurrency and none_deferred:
# small speed up
are_runnable_tasks = (
schedulable_tis
or self._are_premature_tis(unfinished_tis, finished_tis, session)
or changed_tis
)
leaf_task_ids = {t.task_id for t in dag.leaves}
leaf_tis = [ti for ti in tis if ti.task_id in leaf_task_ids if ti.state != TaskInstanceState.REMOVED]
# if all roots finished and at least one failed, the run failed
if not unfinished_tis and any(leaf_ti.state in State.failed_states for leaf_ti in leaf_tis):
self.log.error('Marking run %s failed', self)
self.set_state(DagRunState.FAILED)
if execute_callbacks:
dag.handle_callback(self, success=False, reason='task_failure', session=session)
elif dag.has_on_failure_callback:
callback = DagCallbackRequest(
full_filepath=dag.fileloc,
dag_id=self.dag_id,
run_id=self.run_id,
is_failure_callback=True,
msg='task_failure',
)
# if all leaves succeeded and no unfinished tasks, the run succeeded
elif not unfinished_tis and all(leaf_ti.state in State.success_states for leaf_ti in leaf_tis):
self.log.info('Marking run %s successful', self)
self.set_state(DagRunState.SUCCESS)
if execute_callbacks:
dag.handle_callback(self, success=True, reason='success', session=session)
elif dag.has_on_success_callback:
callback = DagCallbackRequest(
full_filepath=dag.fileloc,
dag_id=self.dag_id,
run_id=self.run_id,
is_failure_callback=False,
msg='success',
)
# if *all tasks* are deadlocked, the run failed
elif (
unfinished_tis
and none_depends_on_past
and none_task_concurrency
and none_deferred
and not are_runnable_tasks
):
self.log.error('Deadlock; marking run %s failed', self)
self.set_state(DagRunState.FAILED)
if execute_callbacks:
dag.handle_callback(self, success=False, reason='all_tasks_deadlocked', session=session)
elif dag.has_on_failure_callback:
callback = DagCallbackRequest(
full_filepath=dag.fileloc,
dag_id=self.dag_id,
run_id=self.run_id,
is_failure_callback=True,
msg='all_tasks_deadlocked',
)
# finally, if the roots aren't done, the dag is still running
else:
self.set_state(DagRunState.RUNNING)
if self._state == DagRunState.FAILED or self._state == DagRunState.SUCCESS:
msg = (
"DagRun Finished: dag_id=%s, execution_date=%s, run_id=%s, "
"run_start_date=%s, run_end_date=%s, run_duration=%s, "
"state=%s, external_trigger=%s, run_type=%s, "
"data_interval_start=%s, data_interval_end=%s, dag_hash=%s"
)
self.log.info(
msg,
self.dag_id,
self.execution_date,
self.run_id,
self.start_date,
self.end_date,
(self.end_date - self.start_date).total_seconds()
if self.start_date and self.end_date
else None,
self._state,
self.external_trigger,
self.run_type,
self.data_interval_start,
self.data_interval_end,
self.dag_hash,
)
session.flush()
self._emit_true_scheduling_delay_stats_for_finished_state(finished_tis)
self._emit_duration_stats_for_finished_state()
session.merge(self)
# We do not flush here for performance reasons(It increases queries count by +20)
return schedulable_tis, callback
@provide_session
def task_instance_scheduling_decisions(self, session: Session = NEW_SESSION) -> TISchedulingDecision:
tis = self.get_task_instances(session=session, state=State.task_states)
self.log.debug("number of tis tasks for %s: %s task(s)", self, len(tis))
def _filter_tis_and_exclude_removed(dag: "DAG", tis: List[TI]) -> Iterable[TI]:
"""Populate ``ti.task`` while excluding those missing one, marking them as REMOVED."""
for ti in tis:
try:
ti.task = dag.get_task(ti.task_id)
except TaskNotFound:
if ti.state != State.REMOVED:
self.log.error("Failed to get task for ti %s. Marking it as removed.", ti)
ti.state = State.REMOVED
session.flush()
else:
yield ti
tis = list(_filter_tis_and_exclude_removed(self.get_dag(), tis))
missing_indexes = self._find_missing_task_indexes(tis, session=session)
if missing_indexes:
self.verify_integrity(missing_indexes=missing_indexes, session=session)
unfinished_tis = [t for t in tis if t.state in State.unfinished]
finished_tis = [t for t in tis if t.state in State.finished]
if unfinished_tis:
schedulable_tis = [ut for ut in unfinished_tis if ut.state in SCHEDULEABLE_STATES]
self.log.debug("number of scheduleable tasks for %s: %s task(s)", self, len(schedulable_tis))
schedulable_tis, changed_tis, expansion_happened = self._get_ready_tis(
schedulable_tis,
finished_tis,
session=session,
)
# During expansion we may change some tis into non-schedulable
# states, so we need to re-compute.
if expansion_happened:
new_unfinished_tis = [t for t in unfinished_tis if t.state in State.unfinished]
finished_tis.extend(t for t in unfinished_tis if t.state in State.finished)
unfinished_tis = new_unfinished_tis
else:
schedulable_tis = []
changed_tis = False
return TISchedulingDecision(
tis=tis,
schedulable_tis=schedulable_tis,
changed_tis=changed_tis,
unfinished_tis=unfinished_tis,
finished_tis=finished_tis,
)
def _get_ready_tis(
self,
schedulable_tis: List[TI],
finished_tis: List[TI],
session: Session,
) -> Tuple[List[TI], bool, bool]:
old_states = {}
ready_tis: List[TI] = []
changed_tis = False
if not schedulable_tis:
return ready_tis, changed_tis, False
# If we expand TIs, we need a new list so that we iterate over them too. (We can't alter
# `schedulable_tis` in place and have the `for` loop pick them up
additional_tis: List[TI] = []
dep_context = DepContext(
flag_upstream_failed=True,
ignore_unmapped_tasks=True, # Ignore this Dep, as we will expand it if we can.
finished_tis=finished_tis,
)
# Check dependencies.
expansion_happened = False
for schedulable in itertools.chain(schedulable_tis, additional_tis):
old_state = schedulable.state
if not schedulable.are_dependencies_met(session=session, dep_context=dep_context):
old_states[schedulable.key] = old_state
continue
# If schedulable is from a mapped task, but not yet expanded, do it
# now. This is called in two places: First and ideally in the mini
# scheduler at the end of LocalTaskJob, and then as an "expansion of
# last resort" in the scheduler to ensure that the mapped task is
# correctly expanded before executed.
if schedulable.map_index < 0 and isinstance(schedulable.task, MappedOperator):
expanded_tis, _ = schedulable.task.expand_mapped_task(self.run_id, session=session)
if expanded_tis:
assert expanded_tis[0] is schedulable
additional_tis.extend(expanded_tis[1:])
expansion_happened = True
if schedulable.state in SCHEDULEABLE_STATES:
ready_tis.append(schedulable)
# Check if any ti changed state
tis_filter = TI.filter_for_tis(old_states)
if tis_filter is not None:
fresh_tis = session.query(TI).filter(tis_filter).all()
changed_tis = any(ti.state != old_states[ti.key] for ti in fresh_tis)
return ready_tis, changed_tis, expansion_happened
def _are_premature_tis(
self,
unfinished_tis: List[TI],
finished_tis: List[TI],
session: Session,
) -> bool:
# there might be runnable tasks that are up for retry and for some reason(retry delay, etc) are
# not ready yet so we set the flags to count them in
for ut in unfinished_tis:
if ut.are_dependencies_met(
dep_context=DepContext(
flag_upstream_failed=True,
ignore_in_retry_period=True,
ignore_in_reschedule_period=True,
finished_tis=finished_tis,
),
session=session,
):
return True
return False
def _emit_true_scheduling_delay_stats_for_finished_state(self, finished_tis: List[TI]) -> None:
"""
This is a helper method to emit the true scheduling delay stats, which is defined as
the time when the first task in DAG starts minus the expected DAG run datetime.
This method will be used in the update_state method when the state of the DagRun
is updated to a completed status (either success or failure). The method will find the first
started task within the DAG and calculate the expected DagRun start time (based on
dag.execution_date & dag.timetable), and minus these two values to get the delay.
The emitted data may contains outlier (e.g. when the first task was cleared, so
the second task's start_date will be used), but we can get rid of the outliers
on the stats side through the dashboards tooling built.
Note, the stat will only be emitted if the DagRun is a scheduler triggered one
(i.e. external_trigger is False).
"""
if self.state == State.RUNNING:
return
if self.external_trigger:
return
if not finished_tis:
return
try:
dag = self.get_dag()
if not dag.timetable.periodic:
# We can't emit this metric if there is no following schedule to calculate from!
return
ordered_tis_by_start_date = [ti for ti in finished_tis if ti.start_date]
ordered_tis_by_start_date.sort(key=lambda ti: ti.start_date, reverse=False)
first_start_date = ordered_tis_by_start_date[0].start_date
if first_start_date:
# TODO: Logically, this should be DagRunInfo.run_after, but the
# information is not stored on a DagRun, only before the actual
# execution on DagModel.next_dagrun_create_after. We should add
# a field on DagRun for this instead of relying on the run
# always happening immediately after the data interval.
data_interval_end = dag.get_run_data_interval(self).end
true_delay = first_start_date - data_interval_end
if true_delay.total_seconds() > 0:
Stats.timing(f'dagrun.{dag.dag_id}.first_task_scheduling_delay', true_delay)
except Exception as e:
self.log.warning(f'Failed to record first_task_scheduling_delay metric:\n{e}')
def _emit_duration_stats_for_finished_state(self):
if self.state == State.RUNNING:
return
if self.start_date is None:
self.log.warning('Failed to record duration of %s: start_date is not set.', self)
return
if self.end_date is None:
self.log.warning('Failed to record duration of %s: end_date is not set.', self)
return
duration = self.end_date - self.start_date
if self.state == State.SUCCESS:
Stats.timing(f'dagrun.duration.success.{self.dag_id}', duration)
elif self.state == State.FAILED:
Stats.timing(f'dagrun.duration.failed.{self.dag_id}', duration)
@provide_session
def verify_integrity(
self,
*,
missing_indexes: Optional[Dict["MappedOperator", Sequence[int]]] = None,
session: Session = NEW_SESSION,
):
"""
Verifies the DagRun by checking for removed tasks or tasks that are not in the
database yet. It will set state to removed or add the task if required.
:missing_indexes: A dictionary of task vs indexes that are missing.
:param session: Sqlalchemy ORM Session
"""
from airflow.settings import task_instance_mutation_hook
# Set for the empty default in airflow.settings -- if it's not set this means it has been changed
hook_is_noop = getattr(task_instance_mutation_hook, 'is_noop', False)
dag = self.get_dag()
task_ids: Set[str] = set()
if missing_indexes:
tis = self.get_task_instances(session=session)
for ti in tis:
task_instance_mutation_hook(ti)
task_ids.add(ti.task_id)
else:
task_ids, missing_indexes = self._check_for_removed_or_restored_tasks(
dag, task_instance_mutation_hook, session=session
)
def task_filter(task: "Operator") -> bool:
return task.task_id not in task_ids and (
self.is_backfill
or task.start_date <= self.execution_date
and (task.end_date is None or self.execution_date <= task.end_date)
)
created_counts: Dict[str, int] = defaultdict(int)
# Get task creator function
task_creator = self._get_task_creator(created_counts, task_instance_mutation_hook, hook_is_noop)
# Create the missing tasks, including mapped tasks
tasks = self._create_missing_tasks(dag, task_creator, task_filter, missing_indexes, session=session)
self._create_task_instances(dag.dag_id, tasks, created_counts, hook_is_noop, session=session)
def _check_for_removed_or_restored_tasks(
self, dag: "DAG", ti_mutation_hook, *, session: Session
) -> Tuple[Set[str], Dict["MappedOperator", Sequence[int]]]:
"""
Check for removed tasks/restored/missing tasks.
:param dag: DAG object corresponding to the dagrun
:param ti_mutation_hook: task_instance_mutation_hook function
:param session: Sqlalchemy ORM Session
:return: List of task_ids in the dagrun and missing task indexes
"""
tis = self.get_task_instances(session=session)
# check for removed or restored tasks
task_ids = set()
existing_indexes: Dict["MappedOperator", List[int]] = defaultdict(list)
expected_indexes: Dict["MappedOperator", Sequence[int]] = defaultdict(list)
for ti in tis:
ti_mutation_hook(ti)
task_ids.add(ti.task_id)
task = None
try:
task = dag.get_task(ti.task_id)
should_restore_task = (task is not None) and ti.state == State.REMOVED
if should_restore_task:
self.log.info("Restoring task '%s' which was previously removed from DAG '%s'", ti, dag)
Stats.incr(f"task_restored_to_dag.{dag.dag_id}", 1, 1)
ti.state = State.NONE
except AirflowException:
if ti.state == State.REMOVED:
pass # ti has already been removed, just ignore it
elif self.state != State.RUNNING and not dag.partial:
self.log.warning("Failed to get task '%s' for dag '%s'. Marking it as removed.", ti, dag)
Stats.incr(f"task_removed_from_dag.{dag.dag_id}", 1, 1)
ti.state = State.REMOVED
continue
if not task.is_mapped:
continue
task = cast("MappedOperator", task)
num_mapped_tis = task.parse_time_mapped_ti_count
# Check if the number of mapped literals has changed and we need to mark this TI as removed
if num_mapped_tis is not None:
if ti.map_index >= num_mapped_tis:
self.log.debug(
"Removing task '%s' as the map_index is longer than the literal mapping list (%s)",
ti,
num_mapped_tis,
)
ti.state = State.REMOVED
elif ti.map_index < 0:
self.log.debug("Removing the unmapped TI '%s' as the mapping can now be performed", ti)
ti.state = State.REMOVED
else:
self.log.info("Restoring mapped task '%s'", ti)
Stats.incr(f"task_restored_to_dag.{dag.dag_id}", 1, 1)
existing_indexes[task].append(ti.map_index)
expected_indexes[task] = range(num_mapped_tis)
else:
# What if it is _now_ dynamically mapped, but wasn't before?
total_length = task.run_time_mapped_ti_count(self.run_id, session=session)
if total_length is None:
# Not all upstreams finished, so we can't tell what should be here. Remove everything.
if ti.map_index >= 0:
self.log.debug(
"Removing the unmapped TI '%s' as the mapping can't be resolved yet", ti
)
ti.state = State.REMOVED
continue
# Upstreams finished, check there aren't any extras
if ti.map_index >= total_length:
self.log.debug(
"Removing task '%s' as the map_index is longer than the resolved mapping list (%d)",
ti,
total_length,
)
ti.state = State.REMOVED
else:
self.log.info("Restoring mapped task '%s'", ti)
Stats.incr(f"task_restored_to_dag.{dag.dag_id}", 1, 1)
existing_indexes[task].append(ti.map_index)
expected_indexes[task] = range(total_length)
# Check if we have some missing indexes to create ti for
missing_indexes: Dict["MappedOperator", Sequence[int]] = defaultdict(list)
for k, v in existing_indexes.items():
missing_indexes.update({k: list(set(expected_indexes[k]).difference(v))})
return task_ids, missing_indexes
def _get_task_creator(
self, created_counts: Dict[str, int], ti_mutation_hook: Callable, hook_is_noop: bool
) -> Callable:
"""
Get the task creator function.
This function also updates the created_counts dictionary with the number of tasks created.
:param created_counts: Dictionary of task_type -> count of created TIs
:param ti_mutation_hook: task_instance_mutation_hook function
:param hook_is_noop: Whether the task_instance_mutation_hook is a noop
"""
if hook_is_noop:
def create_ti_mapping(task: "Operator", indexes: Tuple[int, ...]) -> Generator:
created_counts[task.task_type] += 1
for map_index in indexes:
yield TI.insert_mapping(self.run_id, task, map_index=map_index)
creator = create_ti_mapping
else:
def create_ti(task: "Operator", indexes: Tuple[int, ...]) -> Generator:
for map_index in indexes:
ti = TI(task, run_id=self.run_id, map_index=map_index)
ti_mutation_hook(ti)
created_counts[ti.operator] += 1
yield ti
creator = create_ti
return creator
def _create_missing_tasks(
self,
dag: "DAG",
task_creator: Callable,