Fully support running more than one scheduler concurrently

airflow/api/common/experimental/trigger_dag.py

@@ -89,6 +89,7 @@ def _trigger_dag(
             state=State.RUNNING,
             conf=run_conf,
             external_trigger=True,
+            dag_hash=dag_bag.dags_hash.get(dag_id, None),
         )
 
         triggers.append(trigger)

airflow/cli/cli_parser.py

@@ -1335,7 +1335,7 @@ class GroupCommand(NamedTuple):
         help="Start a scheduler instance",
         func=lazy_load_command('airflow.cli.commands.scheduler_command.scheduler'),
         args=(
-            ARG_DAG_ID_OPT, ARG_SUBDIR, ARG_NUM_RUNS, ARG_DO_PICKLE, ARG_PID, ARG_DAEMON, ARG_STDOUT,
+            ARG_SUBDIR, ARG_NUM_RUNS, ARG_DO_PICKLE, ARG_PID, ARG_DAEMON, ARG_STDOUT,
             ARG_STDERR, ARG_LOG_FILE
         ),
     ),

airflow/cli/commands/scheduler_command.py

@@ -32,7 +32,6 @@ def scheduler(args):
     """Starts Airflow Scheduler"""
     print(settings.HEADER)
     job = SchedulerJob(
-        dag_id=args.dag_id,
         subdir=process_subdir(args.subdir),
         num_runs=args.num_runs,
         do_pickle=args.do_pickle)

airflow/config_templates/config.yml

@@ -559,6 +559,15 @@
       type: string
       example: "connexion,sqlalchemy"
       default: ""
+    - name: max_db_retries
+      description: |
+        Number of times the code should be retried in case of DB Operational Errors.
+        Not all transactions will be retried as it can cause undesired state.
+        Currently it is only used in ``DagFileProcessor.process_file`` to retry ``dagbag.sync_to_db``.
+      version_added: ~
+      type: int
+      example: ~
+      default: "3"
 - name: secrets
   description: ~
   options:
@@ -1572,6 +1581,41 @@
       type: string
       example: ~
       default: "512"
+    - name: use_row_level_locking
+      description: |
+        Should the scheduler issue `SELECT ... FOR UPDATE` in relevant queries.
+        If this is set to False then you should not run more than a single
+        scheduler at once
+      version_added: 2.0.0
+      type: boolean
+      example: ~
+      default: "True"
+    - name: max_dagruns_to_create_per_loop
+      description: |
+        This changes the number of dags that are locked by each scheduler when
+        creating dag runs. One possible reason for setting this lower is if you
+        have huge dags and are running multiple schedules, you won't want one
+        scheduler to do all the work.
+
+        Default: 10
+      example: ~
+      version_added: 2.0.0
+      type: string
+      default: ~
+    - name: max_dagruns_per_loop_to_schedule
+      description: |
+        How many DagRuns should a scheduler examine (and lock) when scheduling
+        and queuing tasks. Increasing this limit will allow more throughput for
+        smaller DAGs but will likely slow down throughput for larger (>500
+        tasks for example) DAGs. Setting this too high when using multiple
+        schedulers could also lead to one scheduler taking all the dag runs
+        leaving no work for the others.
+
+        Default: 20
+      example: ~
+      version_added: 2.0.0
+      type: string
+      default: ~
     - name: statsd_on
       description: |
         Statsd (https://github.com/etsy/statsd) integration settings

airflow/config_templates/default_airflow.cfg

@@ -305,6 +305,11 @@ task_log_reader = task
 # Example: extra_loggers = connexion,sqlalchemy
 extra_loggers =
 
+# Number of times the code should be retried in case of DB Operational Errors.
+# Not all transactions will be retried as it can cause undesired state.
+# Currently it is only used in ``DagFileProcessor.process_file`` to retry ``dagbag.sync_to_db``.
+max_db_retries = 3
+
 [secrets]
 # Full class name of secrets backend to enable (will precede env vars and metastore in search path)
 # Example: backend = airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend
@@ -789,6 +794,29 @@ catchup_by_default = True
 # Set this to 0 for no limit (not advised)
 max_tis_per_query = 512
 
+# Should the scheduler issue `SELECT ... FOR UPDATE` in relevant queries.
+# If this is set to False then you should not run more than a single
+# scheduler at once
+use_row_level_locking = True
+
+# This changes the number of dags that are locked by each scheduler when
+# creating dag runs. One possible reason for setting this lower is if you
+# have huge dags and are running multiple schedules, you won't want one
+# scheduler to do all the work.
+#
+# Default: 10
+# max_dagruns_to_create_per_loop =
+
+# How many DagRuns should a scheduler examine (and lock) when scheduling
+# and queuing tasks. Increasing this limit will allow more throughput for
+# smaller DAGs but will likely slow down throughput for larger (>500
+# tasks for example) DAGs. Setting this too high when using multiple
+# schedulers could also lead to one scheduler taking all the dag runs
+# leaving no work for the others.
+#
+# Default: 20
+# max_dagruns_per_loop_to_schedule =
+
 # Statsd (https://github.com/etsy/statsd) integration settings
 statsd_on = False
 statsd_host = localhost

airflow/executors/base_executor.py

@@ -17,11 +17,12 @@
 """
 Base executor - this is the base class for all the implemented executors.
 """
+import sys
 from collections import OrderedDict
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 from airflow.configuration import conf
-from airflow.models.taskinstance import SimpleTaskInstance, TaskInstance, TaskInstanceKey
+from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
 from airflow.stats import Stats
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import State
@@ -39,8 +40,8 @@
 # Task that is queued. It contains all the information that is
 # needed to run the task.
 #
-# Tuple of: command, priority, queue name, SimpleTaskInstance
-QueuedTaskInstanceType = Tuple[CommandType, int, Optional[str], Union[SimpleTaskInstance, TaskInstance]]
+# Tuple of: command, priority, queue name, TaskInstance
+QueuedTaskInstanceType = Tuple[CommandType, int, Optional[str], TaskInstance]
 
 # Event_buffer dict value type
 # Tuple of: state, info
@@ -72,16 +73,16 @@ def start(self):  # pragma: no cover
         """
 
     def queue_command(self,
-                      simple_task_instance: SimpleTaskInstance,
+                      task_instance: TaskInstance,
                       command: CommandType,
                       priority: int = 1,
                       queue: Optional[str] = None):
         """Queues command to task"""
-        if simple_task_instance.key not in self.queued_tasks and simple_task_instance.key not in self.running:
+        if task_instance.key not in self.queued_tasks and task_instance.key not in self.running:
             self.log.info("Adding to queue: %s", command)
-            self.queued_tasks[simple_task_instance.key] = (command, priority, queue, simple_task_instance)
+            self.queued_tasks[task_instance.key] = (command, priority, queue, task_instance)
         else:
-            self.log.error("could not queue task %s", simple_task_instance.key)
+            self.log.error("could not queue task %s", task_instance.key)
 
     def queue_task_instance(
             self,
@@ -112,7 +113,7 @@ def queue_task_instance(
             pickle_id=pickle_id,
             cfg_path=cfg_path)
         self.queue_command(
-            SimpleTaskInstance(task_instance),
+            task_instance,
             command_list_to_run,
             priority=task_instance.task.priority_weight_total,
             queue=task_instance.task.queue)
@@ -178,13 +179,13 @@ def trigger_tasks(self, open_slots: int) -> None:
         sorted_queue = self.order_queued_tasks_by_priority()
 
         for _ in range(min((open_slots, len(self.queued_tasks)))):
-            key, (command, _, _, simple_ti) = sorted_queue.pop(0)
+            key, (command, _, _, ti) = sorted_queue.pop(0)
             self.queued_tasks.pop(key)
             self.running.add(key)
             self.execute_async(key=key,
                                command=command,
                                queue=None,
-                               executor_config=simple_ti.executor_config)
+                               executor_config=ti.executor_config)
 
     def change_state(self, key: TaskInstanceKey, state: str, info=None) -> None:
         """
@@ -282,6 +283,16 @@ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance
         # Subclasses can do better!
         return tis
 
+    @property
+    def slots_available(self):
+        """
+        Number of new tasks this executor instance can accept
+        """
+        if self.parallelism:
+            return self.parallelism - len(self.running) - len(self.queued_tasks)
+        else:
+            return sys.maxsize
+
     @staticmethod
     def validate_command(command: List[str]) -> None:
         """Check if the command to execute is airflow command"""

airflow/executors/celery_kubernetes_executor.py

@@ -62,17 +62,19 @@ def start(self) -> None:
         self.celery_executor.start()
         self.kubernetes_executor.start()
 
-    def queue_command(self,
-                      simple_task_instance: SimpleTaskInstance,
-                      command: CommandType,
-                      priority: int = 1,
-                      queue: Optional[str] = None):
+    def queue_command(
+        self,
+        task_instance: TaskInstance,
+        command: CommandType,
+        priority: int = 1,
+        queue: Optional[str] = None
+    ):
         """Queues command via celery or kubernetes executor"""
-        executor = self._router(simple_task_instance)
-        self.log.debug("Using executor: %s for %s",
-                       executor.__class__.__name__, simple_task_instance.key
-                       )
-        executor.queue_command(simple_task_instance, command, priority, queue)
+        executor = self._router(task_instance)
+        self.log.debug(
+            "Using executor: %s for %s", executor.__class__.__name__, task_instance.key
+        )
+        executor.queue_command(task_instance, command, priority, queue)
 
     def queue_task_instance(
             self,