-
Notifications
You must be signed in to change notification settings - Fork 17.5k
Add JitteredCronTimetable for deterministic schedule jitter #69705
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| Add ``JitteredCronTimetable``, a ``CronTriggerTimetable`` that offsets each Dag's fire time by a deterministic, per-Dag jitter to spread out Dags sharing a cron expression. |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -18,6 +18,7 @@ | |
|
|
||
| import datetime | ||
| import functools | ||
| import hashlib | ||
| import math | ||
| import operator | ||
| import time | ||
|
|
@@ -598,3 +599,83 @@ def generate_run_id( | |
| suffix = f"{suffix}__{partition_key}" | ||
| suffix = f"{suffix}__{get_random_string()}" | ||
| return run_type.generate_run_id(suffix=suffix) | ||
|
|
||
|
|
||
| class JitteredCronTimetable(CronTriggerTimetable): | ||
| """ | ||
| A :class:`CronTriggerTimetable` that offsets each run by a deterministic, per-DAG jitter. | ||
|
|
||
| Behaves exactly like ``CronTriggerTimetable`` but shifts every fire time by a fixed offset | ||
| derived from ``seed`` and spread across ``[0, max_jitter)``. This avoids the "thundering | ||
| herd" where many DAGs sharing a cron expression (e.g. ``@daily`` -> ``0 0 * * *``) all fire | ||
| at the same instant and overload the scheduler and workers at that boundary. | ||
|
|
||
| The offset is deterministic: the same ``seed`` always maps to the same offset, so runs are | ||
| stable and predictable across scheduler restarts and timetable serialization. DAGs with | ||
| different seeds land in different slots; collisions are possible but harmless (two DAGs | ||
| sharing one minute still beats every DAG firing at once). | ||
|
|
||
| All ``data_interval`` and ``logical_date`` semantics are inherited from | ||
| ``CronTriggerTimetable`` -- only the wall-clock fire time is shifted. | ||
|
|
||
| :param cron: cron string that defines the base schedule. | ||
| :param timezone: timezone used to interpret the cron string. | ||
| :param interval: timedelta that defines the data interval start (see ``CronTriggerTimetable``). | ||
| :param run_immediately: see ``CronTriggerTimetable``. | ||
| :param seed: stable, unique-per-DAG string that determines the offset. The DAG id is a | ||
| natural choice. | ||
| :param max_jitter: upper bound of the jitter window; the offset falls in ``[0, max_jitter)``. | ||
| Keep it smaller than the gap to the next cron boundary so the shift cannot push a run's | ||
| ``logical_date`` across a day or period boundary. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| cron: str, | ||
| *, | ||
| timezone: str | Timezone | FixedTimezone, | ||
| interval: datetime.timedelta | relativedelta = datetime.timedelta(), | ||
| run_immediately: bool | datetime.timedelta = False, | ||
| seed: str, | ||
| max_jitter: datetime.timedelta, | ||
| ) -> None: | ||
| super().__init__(cron, timezone=timezone, interval=interval, run_immediately=run_immediately) | ||
| h = int(hashlib.md5(seed.encode()).hexdigest(), 16) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| self._offset = ( | ||
| datetime.timedelta(seconds=h % int(max_jitter.total_seconds())) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The guard is |
||
| if max_jitter > datetime.timedelta(0) | ||
| else datetime.timedelta(0) | ||
| ) | ||
| self._seed = seed | ||
| self._max_jitter = max_jitter | ||
|
|
||
| def _apply(self, t: DateTime) -> DateTime: | ||
| return t + self._offset | ||
|
|
||
| def _strip(self, t: DateTime) -> DateTime: | ||
| return t - self._offset | ||
|
|
||
| def _get_next(self, current: DateTime) -> DateTime: | ||
| return self._apply(super()._get_next(self._strip(current))) | ||
|
|
||
| def _get_prev(self, current: DateTime) -> DateTime: | ||
| return self._apply(super()._get_prev(self._strip(current))) | ||
|
|
||
| def serialize(self) -> dict[str, Any]: | ||
| data = super().serialize() | ||
| data["seed"] = self._seed | ||
| data["max_jitter"] = self._max_jitter.total_seconds() | ||
| return data | ||
|
|
||
| @classmethod | ||
| def deserialize(cls, data: dict[str, Any]) -> Timetable: | ||
| from airflow.serialization.decoders import decode_interval, decode_run_immediately | ||
|
|
||
| return cls( | ||
| data["expression"], | ||
| timezone=parse_timezone(data["timezone"]), | ||
| interval=decode_interval(data["interval"]), | ||
| run_immediately=decode_run_immediately(data.get("run_immediately", False)), | ||
| seed=data["seed"], | ||
| max_jitter=datetime.timedelta(seconds=data["max_jitter"]), | ||
| ) | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,136 @@ | ||
| # 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. | ||
| from __future__ import annotations | ||
|
|
||
| from datetime import timedelta | ||
|
|
||
| import pendulum | ||
| import pytest | ||
| import time_machine | ||
|
|
||
| from airflow._shared.timezones.timezone import utc | ||
| from airflow.sdk import JitteredCronTimetable as SdkJitteredCronTimetable | ||
| from airflow.serialization.decoders import decode_timetable | ||
| from airflow.serialization.encoders import encode_timetable | ||
| from airflow.timetables.base import DataInterval, TimeRestriction | ||
| from airflow.timetables.trigger import CronTriggerTimetable, JitteredCronTimetable | ||
|
|
||
| CRON = "0 0 * * *" # daily at midnight | ||
| START_DATE = pendulum.DateTime(2021, 9, 4, tzinfo=utc) | ||
| # seed/window known to produce a non-zero offset (used across the behavioural tests) | ||
| SEED = "my_dag" | ||
| MAX_JITTER = timedelta(hours=1) | ||
|
|
||
|
|
||
| def _catchup_run_afters(timetable, count, *, earliest): | ||
| """Return the ``run_after`` of the first ``count`` catchup runs, feeding each back in.""" | ||
| run_afters = [] | ||
| last = None | ||
| for _ in range(count): | ||
| info = timetable.next_dagrun_info( | ||
| last_automated_data_interval=last, | ||
| restriction=TimeRestriction(earliest=earliest, latest=None, catchup=True), | ||
| ) | ||
| assert info is not None | ||
| run_afters.append(info.run_after) | ||
| last = DataInterval.exact(info.run_after) | ||
| return run_afters | ||
|
|
||
|
|
||
| @pytest.mark.parametrize( | ||
| "timezone_name", | ||
| [ | ||
| pytest.param("UTC", id="utc"), | ||
| # Daily runs across the 2026-03-08 spring-forward in this zone: the offset must be | ||
| # applied identically on both sides of the transition. | ||
| pytest.param("America/New_York", id="dst-spring-forward"), | ||
| ], | ||
| ) | ||
| def test_jittered_runs_equal_base_runs_plus_offset(timezone_name): | ||
| """Every jittered run is the plain-cron run shifted by the fixed per-DAG offset. | ||
|
|
||
| Cron/DST correctness is delegated to ``CronTriggerTimetable``; this asserts the offset is | ||
| applied consistently across a catchup sequence (including a DST transition), i.e. the | ||
| strip -> cron -> apply overrides never drift. | ||
| """ | ||
| earliest = pendulum.datetime(2026, 3, 6, tz=timezone_name) | ||
| base = CronTriggerTimetable(CRON, timezone=timezone_name) | ||
| jittered = JitteredCronTimetable(CRON, timezone=timezone_name, seed=SEED, max_jitter=MAX_JITTER) | ||
| offset = jittered._offset | ||
| assert offset > timedelta(0), "seed/window must produce a real shift, else the test is vacuous" | ||
|
|
||
| base_runs = _catchup_run_afters(base, 5, earliest=earliest) | ||
| jittered_runs = _catchup_run_afters(jittered, 5, earliest=earliest) | ||
|
|
||
| assert jittered_runs == [run + offset for run in base_runs] | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("catchup", [True, False]) | ||
| def test_zero_max_jitter_matches_cron_trigger(catchup): | ||
| """A zero window yields a zero offset, so the timetable behaves exactly like its parent.""" | ||
| base = CronTriggerTimetable(CRON, timezone=utc) | ||
| jittered = JitteredCronTimetable(CRON, timezone=utc, seed=SEED, max_jitter=timedelta(0)) | ||
| assert jittered._offset == timedelta(0) | ||
|
|
||
| last = DataInterval.exact(pendulum.DateTime(2022, 7, 26, tzinfo=utc)) | ||
| restriction = TimeRestriction(earliest=START_DATE, latest=None, catchup=catchup) | ||
| # travel so the no-catchup branch (which reads utcnow()) is deterministic | ||
| with time_machine.travel(pendulum.DateTime(2022, 7, 27, 5, 30, tzinfo=utc)): | ||
| assert jittered.next_dagrun_info( | ||
| last_automated_data_interval=last, restriction=restriction | ||
| ) == base.next_dagrun_info(last_automated_data_interval=last, restriction=restriction) | ||
|
|
||
|
|
||
| def test_offsets_are_deterministic_bounded_and_spread(): | ||
| """Same seed -> same offset (stable hash, not process-salted); offsets stay in [0, max_jitter) and spread.""" | ||
| assert ( | ||
| JitteredCronTimetable(CRON, timezone=utc, seed="dag_a", max_jitter=MAX_JITTER)._offset | ||
| == JitteredCronTimetable(CRON, timezone=utc, seed="dag_a", max_jitter=MAX_JITTER)._offset | ||
| ) | ||
|
|
||
| offsets = [ | ||
| JitteredCronTimetable(CRON, timezone=utc, seed=f"dag_{i}", max_jitter=MAX_JITTER)._offset | ||
| for i in range(25) | ||
| ] | ||
| assert all(timedelta(0) <= offset < MAX_JITTER for offset in offsets) | ||
| assert len(set(offsets)) > 1, "distinct seeds should not all collide on one slot" | ||
|
|
||
|
|
||
| def test_serialize_round_trip_preserves_offset(): | ||
| """The core ``serialize``/``deserialize`` round-trips seed + window (as seconds) and the derived offset.""" | ||
| tt = JitteredCronTimetable(CRON, timezone=utc, seed=SEED, max_jitter=MAX_JITTER) | ||
| data = tt.serialize() | ||
| assert data["seed"] == SEED | ||
| assert data["max_jitter"] == 3600.0 # serialized as plain seconds | ||
|
|
||
| restored = JitteredCronTimetable.deserialize(data) | ||
| assert isinstance(restored, JitteredCronTimetable) | ||
| assert restored._seed == tt._seed | ||
| assert restored._max_jitter == tt._max_jitter | ||
| assert restored._offset == tt._offset | ||
|
|
||
|
|
||
| def test_encode_decode_round_trip_across_layers(): | ||
| """SDK class -> encode_timetable (dispatch keyed on the SDK class) -> decode_timetable -> core class.""" | ||
| sdk_tt = SdkJitteredCronTimetable(CRON, timezone="UTC", seed=SEED, max_jitter=MAX_JITTER) | ||
|
|
||
| restored = decode_timetable(encode_timetable(sdk_tt)) | ||
|
|
||
| assert isinstance(restored, JitteredCronTimetable) # rebuilt as the core scheduler-side class | ||
| assert restored._max_jitter == MAX_JITTER | ||
| expected_offset = JitteredCronTimetable(CRON, timezone=utc, seed=SEED, max_jitter=MAX_JITTER)._offset | ||
| assert restored._offset == expected_offset |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
"Only the wall-clock fire time is shifted" reads as if
logical_datestays pinned to the cron boundary, but it doesn't. With the defaultinterval=0,next_dagrun_inforeturnsDagRunInfo.interval(next_start_time, next_start_time)wherenext_start_timeis the jittered time, sological_date(=data_interval.start) moves by the offset too. Themax_jittercaveat further down ("the shift cannot push a run'slogical_dateacross a day or period boundary") already assumes this, so the two statements contradict each other. Worth rewording to say the offset shiftslogical_date/data_intervalas well, or deciding whetherlogical_dateshould stay pinned to the boundary.