Repository: apache/airflow

Issue:
Custom Timetable error with CeleryExecutor
### Apache Airflow version

2.2.3 (latest released)

### What happened

i'm really new to airflow, and i have an error when using Custom Timetable w/ CeleryExecutor

### What you expected to happen

For the custom timetable to be implemented and used by DAG.

error log as follows:

```
[2021-12-23 05:44:30,843: WARNING/ForkPoolWorker-2] Running <TaskInstance: example_cron_trivial_dag.dummy scheduled__2021-12-23T05:44:00+00:00 [queued]> on host 310bbd362d25
[2021-12-23 05:44:30,897: ERROR/ForkPoolWorker-2] Failed to execute task Not a valid timetable: <cron_trivial_timetable.CronTrivialTimetable object at 0x7fc15a0a4100>.
Traceback (most recent call last):
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/executors/celery_executor.py", line 121, in _execute_in_fork
    args.func(args)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/cli/cli_parser.py", line 48, in command
    return func(*args, **kwargs)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/utils/cli.py", line 92, in wrapper
    return f(*args, **kwargs)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/cli/commands/task_command.py", line 298, in task_run
    _run_task_by_selected_method(args, dag, ti)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/cli/commands/task_command.py", line 105, in _run_task_by_selected_method
    _run_task_by_local_task_job(args, ti)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/cli/commands/task_command.py", line 163, in _run_task_by_local_task_job
    run_job.run()
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/jobs/base_job.py", line 245, in run
    self._execute()
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/jobs/local_task_job.py", line 78, in _execute
    self.task_runner = get_task_runner(self)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/task/task_runner/__init__.py", line 63, in get_task_runner
    task_runner = task_runner_class(local_task_job)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/task/task_runner/standard_task_runner.py", line 35, in __init__
    super().__init__(local_task_job)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/task/task_runner/base_task_runner.py", line 48, in __init__
    super().__init__(local_task_job.task_instance)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/utils/log/logging_mixin.py", line 40, in __init__
    self._set_context(context)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/utils/log/logging_mixin.py", line 54, in _set_context
    set_context(self.log, context)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/utils/log/logging_mixin.py", line 178, in set_context
    handler.set_context(value)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/utils/log/file_task_handler.py", line 59, in set_context
    local_loc = self._init_file(ti)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/utils/log/file_task_handler.py", line 264, in _init_file
    relative_path = self._render_filename(ti, ti.try_number)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/utils/log/file_task_handler.py", line 80, in _render_filename
    context = ti.get_template_context()
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1912, in get_template_context
    'next_ds': get_next_ds(),
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1868, in get_next_ds
    execution_date = get_next_execution_date()
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1862, in get_next_execution_date
    next_execution_date = dag.following_schedule(self.execution_date)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/models/dag.py", line 595, in following_schedule
    data_interval = self.infer_automated_data_interval(timezone.coerce_datetime(dttm))
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/models/dag.py", line 678, in infer_automated_data_interval
    raise ValueError(f"Not a valid timetable: {self.timetable!r}")
ValueError: Not a valid timetable: <cron_trivial_timetable.CronTrivialTimetable object at 0x7fc15a0a4100>
[2021-12-23 05:44:30,936: ERROR/ForkPoolWorker-2] Task airflow.executors.celery_executor.execute_command[7c76904d-1b61-4441-b5f0-96ef2ba7c3b7] raised unexpected: AirflowException('Celery command failed on host: 310bbd362d25')
Traceback (most recent call last):
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/celery/app/trace.py", line 451, in trace_task
    R = retval = fun(*args, **kwargs)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/celery/app/trace.py", line 734, in __protected_call__
    return self.run(*args, **kwargs)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/executors/celery_executor.py", line 90, in execute_command
    _execute_in_fork(command_to_exec, celery_task_id)
  File "/opt/rh/rh-python38/root/usr/local/lib/python3.8/site-packages/airflow/executors/celery_executor.py", line 101, in _execute_in_fork
    raise AirflowException('Celery command failed on host: ' + get_hostname())
airflow.exceptions.AirflowException: Celery command failed on host: 310bbd362d25
```

seems like the same problem mentioned here, but by different reasons that i really don't understand
https://github.com/apache/airflow/issues/19578#issuecomment-974654209

### How to reproduce

before 2.2.3 release, i'm having DAG import error w/ custom timetables, and since i only want Next Dag Run on UI be exact time, i will just change the code of CronDataIntervalTimetable by returning DataInterval.exact(end) in next_dagrun_info and infer_manual_data_interval.

here's the new timetable plugin file that i created (simply copied and tweaked from CronDataIntervalTimetable):

```
import datetime
from typing import Any, Dict, Optional, Union

from croniter import CroniterBadCronError, CroniterBadDateError, croniter
from dateutil.relativedelta import relativedelta
from pendulum import DateTime
from pendulum.tz.timezone import Timezone

from airflow.plugins_manager import AirflowPlugin
from airflow.compat.functools import cached_property
from airflow.exceptions import AirflowTimetableInvalid
from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable
from airflow.utils.dates import cron_presets
from airflow.utils.timezone import convert_to_utc, make_aware, make_naive

Delta = Union[datetime.timedelta, relativedelta]

class _TrivialTimetable(Timetable):
    """Basis for timetable implementations that schedule data intervals.

    This kind of timetable classes create periodic data intervals (exact times) from an
    underlying schedule representation (e.g. a cron expression, or a timedelta
    instance), and schedule a DagRun at the end of each interval (exact time).
    """

    def _skip_to_latest(self, earliest: Optional[DateTime]) -> DateTime:
        """Bound the earliest time a run can be scheduled.

        This is called when ``catchup=False``.
        """
        raise NotImplementedError()

    def _align(self, current: DateTime) -> DateTime:
        """Align given time to the scheduled.

        For fixed schedules (e.g. every midnight); this finds the next time that
        aligns to the declared time, if the given time does not align. If the
        schedule is not fixed (e.g. every hour), the given time is returned.
        """
        raise NotImplementedError()

    def _get_next(self, current: DateTime) -> DateTime:
        """Get the first schedule after the current time."""
        raise NotImplementedError()

    def _get_prev(self, current: DateTime) -> DateTime:
        """Get the last schedule before the current time."""
        raise NotImplementedError()

    def next_dagrun_info(
        self,
        *,
        last_automated_data_interval: Optional[DataInterval],
        restriction: TimeRestriction,
    ) -> Optional[DagRunInfo]:
        earliest = restriction.earliest
        if not restriction.catchup:
            earliest = self._skip_to_latest(earliest)
        if last_automated_data_interval is None:
            # First run; schedule the run at the first available time matching
            # the schedule, and retrospectively create a data interval for it.
            if earliest is None:
                return None
            start = self._align(earliest)
        else:
            # There's a previous run. Create a data interval starting from when
            # the end of the previous interval.
            start = last_automated_data_interval.end
        if restriction.latest is not None and start > restriction.latest:
            return None
        end = self._get_next(start)
        return DagRunInfo.exact(end)

def _is_schedule_fixed(expression: str) -> bool:
    """Figures out if the schedule has a fixed time (e.g. 3 AM every day).

    :return: True if the schedule has a fixed time, False if not.

    Detection is done by "peeking" the next two cron trigger time; if the
    two times have the same minute and hour value, the schedule is fixed,
    and we *don't* need to perform the DST fix.

    This assumes DST happens on whole minute changes (e.g. 12:59 -> 12:00).
    """
    cron = croniter(expression)
    next_a = cron.get_next(datetime.datetime)
    next_b = cron.get_next(datetime.datetime)
    return next_b.minute == next_a.minute and next_b.hour == next_a.hour

class CronTrivialTimetable(_TrivialTimetable):
    """Timetable that schedules data intervals with a cron expression.

    This corresponds to ``schedule_interval=<cron>``, where ``<cron>`` is either
    a five/six-segment representation, or one of ``cron_presets``.

    The implementation extends on croniter to add timezone awareness. This is
    because crontier works only with naive timestamps, and cannot consider DST
    when determining the next/previous time.

    Don't pass ``@once`` in here; use ``OnceTimetable`` instead.
    """

    def __init__(self, cron: str, timezone: Timezone) -> None:
        self._expression = cron_presets.get(cron, cron)
        self._timezone = timezone

    @classmethod
    def deserialize(cls, data: Dict[str, Any]) -> "Timetable":
        from airflow.serialization.serialized_objects import decode_timezone

        return cls(data["expression"], decode_timezone(data["timezone"]))

    def __eq__(self, other: Any) -> bool:
        """Both expression and timezone should match.

        This is only for testing purposes and should not be relied on otherwise.
        """
        if not isinstance(other, CronTrivialTimetable):
            return NotImplemented
        return self._expression == other._expression and self._timezone == other._timezone

    @property
    def summary(self) -> str:
        return self._expression
   
    def serialize(self) -> Dict[str, Any]:
        from airflow.serialization.serialized_objects import encode_timezone

        return {"expression": self._expression, "timezone": encode_timezone(self._timezone)}
   
    def validate(self) -> None:
        try:
            croniter(self._expression)
        except (CroniterBadCronError, CroniterBadDateError) as e:
            raise AirflowTimetableInvalid(str(e))
   
    @cached_property
    def _should_fix_dst(self) -> bool:
        # This is lazy so instantiating a schedule does not immediately raise
        # an exception. Validity is checked with validate() during DAG-bagging.
        return not _is_schedule_fixed(self._expression)
   
    def _get_next(self, current: DateTime) -> DateTime:
        """Get the first schedule after specified time, with DST fixed."""
        naive = make_naive(current, self._timezone)
        cron = croniter(self._expression, start_time=naive)
        scheduled = cron.get_next(datetime.datetime)
        if not self._should_fix_dst:
            return convert_to_utc(make_aware(scheduled, self._timezone))
        delta = scheduled - naive
        return convert_to_utc(current.in_timezone(self._timezone) + delta)

    def _get_prev(self, current: DateTime) -> DateTime:
        """Get the first schedule before specified time, with DST fixed."""
        naive = make_naive(current, self._timezone)
        cron = croniter(self._expression, start_time=naive)
        scheduled = cron.get_prev(datetime.datetime)
        if not self._should_fix_dst:
            return convert_to_utc(make_aware(scheduled, self._timezone))
        delta = naive - scheduled
        return convert_to_utc(current.in_timezone(self._timezone) - delta)

    def _align(self, current: DateTime) -> DateTime:
        """Get the next scheduled time.

        This is ``current + interval``, unless ``current`` is first interval,
        then ``current`` is returned.
        """
        next_time = self._get_next(current)
        if self._get_prev(next_time) != current:
            return next_time
        return current

    def _skip_to_latest(self, earliest: Optional[DateTime]) -> DateTime:
        """Bound the earliest time a run can be scheduled.

        The logic is that we move start_date up until one period before, so the
        current time is AFTER the period end, and the job can be created...
        """
        current_time = DateTime.utcnow()
        next_start = self._get_next(current_time)
        last_start = self._get_prev(current_time)
        if next_start == current_time:
            new_start = last_start
        elif next_start > current_time:
            new_start = self._get_prev(last_start)
        else:
            raise AssertionError("next schedule shouldn't be earlier")
        if earliest is None:
            return new_start
        return max(new_start, earliest)

    def infer_manual_data_interval(self, *, run_after: DateTime) -> DataInterval:
        # Get the last complete period before run_after, e.g. if a DAG run is
        # scheduled at each midnight, the data interval of a manually triggered
        # run at 1am 25th is between 0am 24th and 0am 25th.
        end = self._get_prev(self._align(run_after))
        return DataInterval.exact(end)

class CronTrivialTimetablePlugin(AirflowPlugin):
    name = "cron_trivial_timetable_plugin"
    timetables = [CronTrivialTimetable]
```

the dag file i'm using/testing:

```
import pendulum
from datetime import datetime, timedelta
from typing import Dict

from airflow.decorators import task
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.dummy import DummyOperator

from cron_trivial_timetable import CronTrivialTimetable

with DAG(
    dag_id="example_cron_trivial_dag",
    start_date=datetime(2021,11,14,12,0,tzinfo=pendulum.timezone('Asia/Tokyo')),
    max_active_runs=1,
    timetable=CronTrivialTimetable('*/2 * * * *', pendulum.timezone('Asia/Tokyo')),
    default_args={
        'owner': '********',
        'depends_on_past': False,
        'email': [**************],
        'email_on_failure': False,
        'email_on_retry': False,
        'retries': 0,
        'retry_delay': timedelta(seconds=1),
        'end_date': datetime(2101, 1, 1),
    },
    tags=['testing'],
    catchup=False
) as dag:

    dummy = BashOperator(task_id='dummy', queue='daybatch', bash_command="date")
```

### Operating System

CentOS-7

### Versions of Apache Airflow Providers

_No response_

### Deployment

Docker-Compose

### Deployment details

my scheduler & workers are in different dockers, hence i'm using CeleryExecutor

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:

Resolving patch:
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1845,14 +1845,14 @@ def get_prev_start_date_success() -> Optional[pendulum.DateTime]:
 
         @cache
         def get_yesterday_ds() -> str:
-            return (self.execution_date - timedelta(1)).strftime('%Y-%m-%d')
+            return (logical_date - timedelta(1)).strftime('%Y-%m-%d')
 
         def get_yesterday_ds_nodash() -> str:
             return get_yesterday_ds().replace('-', '')
 
         @cache
         def get_tomorrow_ds() -> str:
-            return (self.execution_date + timedelta(1)).strftime('%Y-%m-%d')
+            return (logical_date + timedelta(1)).strftime('%Y-%m-%d')
 
         def get_tomorrow_ds_nodash() -> str:
             return get_tomorrow_ds().replace('-', '')
@@ -1860,18 +1860,15 @@ def get_tomorrow_ds_nodash() -> str:
         @cache
         def get_next_execution_date() -> Optional[pendulum.DateTime]:
             # For manually triggered dagruns that aren't run on a schedule,
-            # next/previous execution dates don't make sense, and should be set
+            # the "next" execution date doesn't make sense, and should be set
             # to execution date for consistency with how execution_date is set
             # for manually triggered tasks, i.e. triggered_date == execution_date.
             if dag_run.external_trigger:
-                next_execution_date = dag_run.execution_date
-            else:
-                with warnings.catch_warnings():
-                    warnings.simplefilter("ignore", DeprecationWarning)
-                    next_execution_date = dag.following_schedule(self.execution_date)
-            if next_execution_date is None:
+                return logical_date
+            next_info = dag.next_dagrun_info(data_interval, restricted=False)
+            if next_info is None:
                 return None
-            return timezone.coerce_datetime(next_execution_date)
+            return timezone.coerce_datetime(next_info.logical_date)
 
         def get_next_ds() -> Optional[str]:
             execution_date = get_next_execution_date()
@@ -1887,11 +1884,15 @@ def get_next_ds_nodash() -> Optional[str]:
 
         @cache
         def get_prev_execution_date():
+            # For manually triggered dagruns that aren't run on a schedule,
+            # the "previous" execution date doesn't make sense, and should be set
+            # to execution date for consistency with how execution_date is set
+            # for manually triggered tasks, i.e. triggered_date == execution_date.
             if dag_run.external_trigger:
-                return timezone.coerce_datetime(self.execution_date)
+                return logical_date
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore", DeprecationWarning)
-                return dag.previous_schedule(self.execution_date)
+                return dag.previous_schedule(logical_date)
 
         @cache
         def get_prev_ds() -> Optional[str]:


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -699,7 +699,15 @@ def create_task_instance(dag_maker, create_dummy_dag):
     Uses ``create_dummy_dag`` to create the dag structure.
     """
 
-    def maker(execution_date=None, dagrun_state=None, state=None, run_id=None, run_type=None, **kwargs):
+    def maker(
+        execution_date=None,
+        dagrun_state=None,
+        state=None,
+        run_id=None,
+        run_type=None,
+        data_interval=None,
+        **kwargs,
+    ):
         if execution_date is None:
             from airflow.utils import timezone
 
@@ -711,6 +719,8 @@ def maker(execution_date=None, dagrun_state=None, state=None, run_id=None, run_t
             dagrun_kwargs["run_id"] = run_id
         if run_type is not None:
             dagrun_kwargs["run_type"] = run_type
+        if data_interval is not None:
+            dagrun_kwargs["data_interval"] = data_interval
         dagrun = dag_maker.create_dagrun(**dagrun_kwargs)
         (ti,) = dagrun.task_instances
         ti.state = state
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -30,6 +30,7 @@
 from freezegun import freeze_time
 
 from airflow import models, settings
+from airflow.example_dags.plugins.workday import AfterWorkdayTimetable
 from airflow.exceptions import (
     AirflowException,
     AirflowFailException,
@@ -1634,6 +1635,30 @@ def test_template_with_json_variable_missing(self, create_task_instance):
         with pytest.raises(KeyError):
             ti.task.render_template('{{ var.json.get("missing_variable") }}', context)
 
+    def test_tempalte_with_custom_timetable_deprecated_context(self, create_task_instance):
+        ti = create_task_instance(
+            start_date=DEFAULT_DATE,
+            timetable=AfterWorkdayTimetable(),
+            run_type=DagRunType.SCHEDULED,
+            execution_date=timezone.datetime(2021, 9, 6),
+            data_interval=(timezone.datetime(2021, 9, 6), timezone.datetime(2021, 9, 7)),
+        )
+        context = ti.get_template_context()
+        with pytest.deprecated_call():
+            assert context["execution_date"] == pendulum.DateTime(2021, 9, 6, tzinfo=timezone.TIMEZONE)
+        with pytest.deprecated_call():
+            assert context["next_ds"] == "2021-09-07"
+        with pytest.deprecated_call():
+            assert context["next_ds_nodash"] == "20210907"
+        with pytest.deprecated_call():
+            assert context["next_execution_date"] == pendulum.DateTime(2021, 9, 7, tzinfo=timezone.TIMEZONE)
+        with pytest.deprecated_call():
+            assert context["prev_ds"] is None, "Does not make sense for custom timetable"
+        with pytest.deprecated_call():
+            assert context["prev_ds_nodash"] is None, "Does not make sense for custom timetable"
+        with pytest.deprecated_call():
+            assert context["prev_execution_date"] is None, "Does not make sense for custom timetable"
+
     def test_execute_callback(self, create_task_instance):
         called = False