Repository: TobikoData/sqlmesh
Pull request #2822 (2024-06-26 00:37:01)

feat: signals for the built-in scheduler
This is just a sketch for #2297. I'll add tests and documentation later on. I'd like feedback on the approach, namely:

1. how should the user provide `signal_factory`?
    1. should they provide a function when building `Context` (like here) or should there be a different mechanism?
1. should waiting occur where it is in this PR or elsewhere? 
1. what about feedback in the console? should there be specific steps output when waiting on signals?
1. should the interface implemented by users be polling or simply blocking (like in the PR)
    1. it's simpler for users to implement a single polling method (i.e.: check the signal once) and the looping is done once
    1. it removes the ability for users to have custom polling / looping (both a good and bad thing)
    1. a poll-based approach makes it easier to give feedback in the console (if necessary)
    1. could also allow a global `poll_interval` and `timeout` configuration (used for every signal)

An alternate interface would be something like this:

```python
class Signal():
  def poll_ready() -> bool:
    """Return true if the signal is ready, false otherwise"
```

And then the looping would be done in the scheduler, e.g.:

```python
poll_timeout = ... # configurable polling timeout
poll_interval = ... # configurable polling interval

for signal in signals:
  start_polling_ts = now()
  while not signal.poll():
    if now() - start_polling_ts > poll_timeout:
      raise TimeoutError
    time.sleep(poll_interval)
```

Ideally, we'd release the thread to go do some other stuff in the meantime, say because there are 2 threads, 3 tasks available, but only 2 of them have a signal; ideally, the 2 signals wouldn't block the third task. It's an optimization, but maybe it's worth it?

izeigerman: Hey @plaflamme , thanks a lot for submitting this PR and kicking off the conversation.

Part of me wonders whether the builtin SQLMesh scheduler is fundamentally too naive and simplistic to support something like external signals.

Here's one specific example that bothers me. Imagine we have some Model A which has a set of upstream dependencies U and a set of downstream dependencies D. In Airflow, when the model A blocks on external signals, obviously all its downstream dependencies (D) are blocked as well.  At the same time, the upstream set U is free to refresh at its own cadence, unimpacted by the fact that Model A is blocked.

Now, in the case of the built-in scheduler, the blocked Model A will block the entire `sqlmesh run` process until the signal requirement is met or the timeout occurs. Which means that upstream models U (as well as models that have no connection to Model A whatsoever) won't be updated until the signal unblocks, the process exits, and subsequently runs again.

The argument can be made that this sort of issue can be mitigated by setting a reasonable timeout. But what is reasonable? What if the external source is experiencing an outage and is delayed by hour(s)? What if the model's cron doesn't match the source update cadence and the source is consistently late? Should we set the timeout to some low value and just fail the `sqlmesh run` execution during this time? That could result in a subpar experience. Or should the scheduler be smarter and run continuously instead of in a fire-and-forget fashion?

TL;DR I believe significant improvements might be required to the built-in scheduler itself before we can consider supporting external signals for it.
tobymao: I have another idea that we can use to create "signals" for the built in scheduler. Build in scheduler "signals" will simply be boolean blockers. If they return True, then they process, else they stop. It's up to your cron to retry, let's say every 5 minutes. 

Imagine U -> (S)ignal -> A.

Run, U process, S blocks. Exit.
Run, U process, S Runs. A runs.

It's a very simple mechanism that allows you to choose whether or not the dag continues. And it's simple to maintain because you just do sqlmesh run at your minimum frequency.
izeigerman: I kind of like the @tobymao's idea. This feels like a natural extension of how we treat cron expressions. 

SQLMesh doesn't block and wait until the designated time to run; it just skips the execution and picks it up again the next time it runs.
plaflamme: @tobymao That does simplify things. One thing that comes to mind though is that the dependencies aren't between models, but between model **intervals**. So, I believe you can get into a "permanently blocked" situation if we're not blocking.

Imagine `U` -> `E`xternal process -> `S`ignal -> `A`.

Where `E` is some external process that waits for intervals in `U` to show up and produce its own output.
And `S` is configured to look for intervals in `E`, say by using `@start_ds` and `@end_ds` in its definition.

Now imagine a stream of `sqlmesh run`s:

| Run | Intervals to run | Signals to check | Signals Ready | Intervals Ready |
| --- |---|---|---|---|
| 1 | `U1`, `A1` | `S1` | | `U1` |
| 2 | `U2`, `A1-2`| `S1-2` | `S1` | `U2` |
| 3 | `U3`, `A1-3`| `S1-3` | `S1-2` |`U3` |

So because on each `sqlmesh run`, the intervals for `A` "grow" at the same rate as its upstream dependency, the signal can never say it is ready, because it will always depend on the interval being run now.

With the blocking approach, this is resolved eventually and `A` will run at some point (baring errors and timeouts).

Maybe this is an edge case, but it doesn't seem far fetched to want to coordinate sqlmesh with external systems like this. Especially during a migration or similar.

Presumably, the signal could return a list of intervals that are ready instead of just a boolean value, but that seems pretty challenging to integrate in practice. Or maybe the scheduler can be smart enough to try different batches somehow.

Another thing that comes to mind is that maybe the user can solve this using `lookback`. That is, if `A` has `lookback 1` then perhaps its `S`ignal can also ignore the absence of the latest interval since it knows it will run it again later. Perhaps this can be provided as a workaround for users with this situation.

FWIW: I don't think this situation applies to us, so I'd be happy to move forward regardless.
plaflamme: @tobymao all checks are passing please let me know if this needs anything else

Patch:
diff --git a/sqlmesh/core/config/scheduler.py b/sqlmesh/core/config/scheduler.py
index 093bdd4d..14159f65 100644
--- a/sqlmesh/core/config/scheduler.py
+++ b/sqlmesh/core/config/scheduler.py
@@ -136,6 +136,7 @@ class BuiltInSchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig)
             backfill_concurrent_tasks=context.concurrent_tasks,
             console=context.console,
             notification_target_manager=context.notification_target_manager,
+            signal_factory=context._signal_factory,
         )
 
     def get_default_catalog(self, context: GenericContext) -> t.Optional[str]:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index b66e9a5d..5c0fd8c6 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -80,7 +80,7 @@ from sqlmesh.core.notification_target import (
 )
 from sqlmesh.core.plan import Plan, PlanBuilder
 from sqlmesh.core.reference import ReferenceGraph
-from sqlmesh.core.scheduler import Scheduler
+from sqlmesh.core.scheduler import Scheduler, SignalFactory
 from sqlmesh.core.schema_loader import create_external_models_file
 from sqlmesh.core.selector import Selector
 from sqlmesh.core.snapshot import (
@@ -308,6 +308,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         load: bool = True,
         console: t.Optional[Console] = None,
         users: t.Optional[t.List[User]] = None,
+        signal_factory: t.Optional[SignalFactory] = None,
     ):
         self.configs = (
             config if isinstance(config, dict) else load_configs(config, self.CONFIG_TYPE, paths)
@@ -354,6 +355,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         )
 
         self._snapshot_evaluator: t.Optional[SnapshotEvaluator] = None
+        self._signal_factory = signal_factory
 
         self._provided_state_sync: t.Optional[StateSync] = state_sync
         self._state_sync: t.Optional[StateSync] = None
@@ -474,6 +476,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             max_workers=self.concurrent_tasks,
             console=self.console,
             notification_target_manager=self.notification_target_manager,
+            signal_factory=self._signal_factory,
         )
 
     @property
diff --git a/sqlmesh/core/plan/evaluator.py b/sqlmesh/core/plan/evaluator.py
index 0e5e98e0..5d408154 100644
--- a/sqlmesh/core/plan/evaluator.py
+++ b/sqlmesh/core/plan/evaluator.py
@@ -26,7 +26,7 @@ from sqlmesh.core.notification_target import (
     NotificationTargetManager,
 )
 from sqlmesh.core.plan.definition import Plan
-from sqlmesh.core.scheduler import Scheduler
+from sqlmesh.core.scheduler import Scheduler, SignalFactory
 from sqlmesh.core.snapshot import DeployabilityIndex, Snapshot, SnapshotEvaluator
 from sqlmesh.core.state_sync import StateSync
 from sqlmesh.core.state_sync.base import PromotionResult
@@ -65,6 +65,7 @@ class BuiltInPlanEvaluator(PlanEvaluator):
         backfill_concurrent_tasks: int = 1,
         console: t.Optional[Console] = None,
         notification_target_manager: t.Optional[NotificationTargetManager] = None,
+        signal_factory: t.Optional[SignalFactory] = None,
     ):
         self.state_sync = state_sync
         self.snapshot_evaluator = snapshot_evaluator
@@ -72,6 +73,7 @@ class BuiltInPlanEvaluator(PlanEvaluator):
         self.backfill_concurrent_tasks = backfill_concurrent_tasks
         self.console = console or get_console()
         self.notification_target_manager = notification_target_manager
+        self.signal_factory = signal_factory
 
     def evaluate(
         self,
@@ -159,6 +161,7 @@ class BuiltInPlanEvaluator(PlanEvaluator):
             max_workers=self.backfill_concurrent_tasks,
             console=self.console,
             notification_target_manager=self.notification_target_manager,
+            signal_factory=self.signal_factory,
         )
         is_run_successful = scheduler.run(
             plan.environment_naming_info,
diff --git a/sqlmesh/core/scheduler.py b/sqlmesh/core/scheduler.py
index f0ac453a..6364b842 100644
--- a/sqlmesh/core/scheduler.py
+++ b/sqlmesh/core/scheduler.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import abc
 import logging
 import traceback
 import typing as t
@@ -30,6 +31,7 @@ from sqlmesh.utils.date import (
     now,
     now_timestamp,
     to_datetime,
+    to_timestamp,
     validate_date_range,
 )
 from sqlmesh.utils.errors import AuditError, CircuitBreakerError, SQLMeshError
@@ -43,6 +45,39 @@ SnapshotToBatches = t.Dict[Snapshot, Batch]
 SchedulingUnit = t.Tuple[str, t.Tuple[Interval, int]]
 
 
+class Signal(abc.ABC):
+    @abc.abstractmethod
+    def check_intervals(self, batch: Batch) -> t.Union[bool, Batch]:
+        """Returns which intervals are ready from a list of scheduled intervals.
+
+        When SQLMesh wishes to execute a batch of intervals, say between `a` and `d`, then
+        the `batch` parameter will contain each individual interval within this batch,
+        i.e.: `[a,b),[b,c),[c,d)`.
+
+        This function may return `True` to indicate that the whole batch is ready,
+        `False` to indicate none of the batch's intervals are ready, or a list of
+        intervals (a batch) to indicate exactly which ones are ready.
+
+        When returning a batch, the function is expected to return a subset of
+        the `batch` parameter, e.g.: `[a,b),[b,c)`. Note that it may return
+        gaps, e.g.: `[a,b),[c,d)`, but it may not alter the bounds of any of the
+        intervals.
+
+        The interface allows an implementation to check batches of intervals without
+        having to actually compute individual intervals itself.
+
+        Args:
+            batch: the list of intervals that are missing and scheduled to run.
+
+        Returns:
+            Either `True` to indicate all intervals are ready, `False` to indicate none are
+            ready or a list of intervals to indicate exactly which ones are ready.
+        """
+
+
+SignalFactory = t.Callable[[t.Dict[str, t.Union[str, int, float, bool]]], Signal]
+
+
 class Scheduler:
     """Schedules and manages the evaluation of snapshots.
 
@@ -58,6 +93,7 @@ class Scheduler:
         state_sync: The state sync to pull saved snapshots.
         max_workers: The maximum number of parallel queries to run.
         console: The rich instance used for printing scheduling information.
+        signal_factory: A factory method for building Signal instances from model signal configuration.
     """
 
     def __init__(
@@ -69,6 +105,7 @@ class Scheduler:
         max_workers: int = 1,
         console: t.Optional[Console] = None,
         notification_target_manager: t.Optional[NotificationTargetManager] = None,
+        signal_factory: t.Optional[SignalFactory] = None,
     ):
         self.state_sync = state_sync
         self.snapshots = {s.snapshot_id: s for s in snapshots}
@@ -80,6 +117,7 @@ class Scheduler:
         self.notification_target_manager = (
             notification_target_manager or NotificationTargetManager()
         )
+        self.signal_factory = signal_factory
 
     def batches(
         self,
@@ -131,6 +169,7 @@ class Scheduler:
             restatements=restatements,
             ignore_cron=ignore_cron,
             end_bounded=end_bounded,
+            signal_factory=self.signal_factory,
         )
 
     def evaluate(
@@ -393,6 +432,7 @@ def compute_interval_params(
     restatements: t.Optional[t.Dict[SnapshotId, SnapshotInterval]] = None,
     ignore_cron: bool = False,
     end_bounded: bool = False,
+    signal_factory: t.Optional[SignalFactory] = None,
 ) -> SnapshotToBatches:
     """Find the optimal date interval paramaters based on what needs processing and maximal batch size.
 
@@ -431,6 +471,15 @@ def compute_interval_params(
         ignore_cron=ignore_cron,
         end_bounded=end_bounded,
     ).items():
+        if signal_factory:
+            for signal in snapshot.model.render_signals(
+                start=start, end=end, execution_time=execution_time
+            ):
+                intervals = _check_ready_intervals(
+                    signal=signal_factory(signal),
+                    intervals=intervals,
+                )
+
         batches = []
         batch_size = snapshot.node.batch_size
         next_batch: t.List[t.Tuple[int, int]] = []
@@ -465,3 +514,53 @@ def _resolve_one_snapshot_per_version(
                 snapshot_per_version[key] = snapshot
 
     return snapshot_per_version
+
+
+def _contiguous_intervals(
+    intervals: t.List[SnapshotInterval],
+) -> t.List[t.List[SnapshotInterval]]:
+    """Given a list of intervals with gaps, returns a list of sequences of contiguous intervals."""
+    contiguous_intervals = []
+    current_batch: t.List[SnapshotInterval] = []
+    for interval in intervals:
+        if len(current_batch) == 0 or interval[0] == current_batch[-1][-1]:
+            current_batch.append(interval)
+        else:
+            contiguous_intervals.append(current_batch)
+            current_batch = [interval]
+
+    if len(current_batch) > 0:
+        contiguous_intervals.append(current_batch)
+
+    return contiguous_intervals
+
+
+def _check_ready_intervals(
+    signal: Signal,
+    intervals: t.List[SnapshotInterval],
+) -> t.List[SnapshotInterval]:
+    """Returns a list of intervals that are considered ready by the provided signal.
+
+    Note that this will handle gaps in the provided intervals. The returned intervals
+    may introduce new gaps.
+    """
+    checked_intervals = []
+    for interval_batch in _contiguous_intervals(intervals):
+        batch = [(to_datetime(start), to_datetime(end)) for start, end in interval_batch]
+
+        ready_intervals = signal.check_intervals(batch=batch)
+        if isinstance(ready_intervals, bool):
+            if not ready_intervals:
+                batch = []
+        elif isinstance(ready_intervals, list):
+            for i in ready_intervals:
+                if i not in batch:
+                    raise RuntimeError(f"Signal returned unknown interval {i}")
+            batch = ready_intervals
+        else:
+            raise ValueError(
+                f"unexpected return value from signal, expected bool | list, got {type(ready_intervals)}"
+            )
+
+        checked_intervals.extend([(to_timestamp(start), to_timestamp(end)) for start, end in batch])
+    return checked_intervals
diff --git a/sqlmesh/schedulers/airflow/plan.py b/sqlmesh/schedulers/airflow/plan.py
index 1a2dc553..e5941f54 100644
--- a/sqlmesh/schedulers/airflow/plan.py
+++ b/sqlmesh/schedulers/airflow/plan.py
@@ -134,6 +134,7 @@ def create_plan_dag_spec(
             deployability_index=deployability_index_for_evaluation,
             restatements=restatements,
             end_bounded=request.end_bounded,
+            signal_factory=None,
         )
     else:
         backfill_batches = {}


Test patch:
diff --git a/tests/core/test_scheduler.py b/tests/core/test_scheduler.py
index c68dbd1b..a0ef5a89 100644
--- a/tests/core/test_scheduler.py
+++ b/tests/core/test_scheduler.py
@@ -458,3 +458,91 @@ def test_external_model_audit(mocker, make_snapshot):
     )
 
     spy.assert_called_once()
+
+
+def test_contiguous_intervals():
+    from sqlmesh.core.scheduler import _contiguous_intervals as ci
+
+    assert ci([]) == []
+    assert ci([(0, 1)]) == [[(0, 1)]]
+    assert ci([(0, 1), (1, 2), (2, 3)]) == [[(0, 1), (1, 2), (2, 3)]]
+    assert ci([(0, 1), (3, 4), (4, 5), (6, 7)]) == [
+        [(0, 1)],
+        [(3, 4), (4, 5)],
+        [(6, 7)],
+    ]
+
+
+def test_check_ready_intervals(mocker: MockerFixture):
+    from sqlmesh.core.scheduler import _check_ready_intervals, Interval
+
+    def const_signal(const):
+        signal_mock = mocker.Mock()
+        signal_mock.check_intervals = mocker.MagicMock(return_value=const)
+        return signal_mock
+
+    def assert_always_signal(intervals):
+        _check_ready_intervals(const_signal(True), intervals) == intervals
+
+    assert_always_signal([])
+    assert_always_signal([(0, 1)])
+    assert_always_signal([(0, 1), (1, 2)])
+    assert_always_signal([(0, 1), (2, 3)])
+
+    def assert_never_signal(intervals):
+        _check_ready_intervals(const_signal(False), intervals) == []
+
+    assert_never_signal([])
+    assert_never_signal([(0, 1)])
+    assert_never_signal([(0, 1), (1, 2)])
+    assert_never_signal([(0, 1), (2, 3)])
+
+    def to_intervals(values: t.List[t.Tuple[int, int]]) -> t.List[Interval]:
+        return [(to_datetime(s), to_datetime(e)) for s, e in values]
+
+    def assert_check_intervals(
+        intervals: t.List[t.Tuple[int, int]],
+        ready: t.List[t.List[t.Tuple[int, int]]],
+        expected: t.List[t.Tuple[int, int]],
+    ):
+        signal_mock = mocker.Mock()
+        signal_mock.check_intervals = mocker.MagicMock(side_effect=[to_intervals(r) for r in ready])
+        _check_ready_intervals(signal_mock, intervals) == expected
+
+    assert_check_intervals([], [], [])
+    assert_check_intervals([(0, 1)], [[]], [])
+    assert_check_intervals(
+        [(0, 1)],
+        [[(0, 1)]],
+        [(0, 1)],
+    )
+    assert_check_intervals(
+        [(0, 1), (1, 2)],
+        [[(0, 1)]],
+        [(0, 1)],
+    )
+    assert_check_intervals(
+        [(0, 1), (1, 2)],
+        [[(1, 2)]],
+        [(1, 2)],
+    )
+    assert_check_intervals(
+        [(0, 1), (1, 2)],
+        [[(0, 1), (1, 2)]],
+        [(0, 1), (1, 2)],
+    )
+    assert_check_intervals(
+        [(0, 1), (1, 2), (3, 4)],
+        [[], []],
+        [],
+    )
+    assert_check_intervals(
+        [(0, 1), (1, 2), (3, 4)],
+        [[(0, 1)], []],
+        [(0, 1)],
+    )
+    assert_check_intervals(
+        [(0, 1), (1, 2), (3, 4)],
+        [[(0, 1)], [(3, 4)]],
+        [(0, 1), (3, 4)],
+    )