Repository: apache/airflow
File: airflow-core/src/airflow/jobs/scheduler_job_runner.py (Python)

Code before:
            .where(TI.state.in_(EXECUTION_STATES))
            .group_by(TI.task_id, TI.run_id, TI.dag_id)
        )
        for dag_id, task_id, run_id, c in query:
            self.dag_run_active_tasks_map[dag_id, run_id] += c
            self.task_concurrency_map[(dag_id, task_id)] += c
            self.task_dagrun_concurrency_map[(dag_id, run_id, task_id)] += c


def _is_parent_process() -> bool:
    """
    Whether this is a parent process.

    Return True if the current process is the parent process.
    False if the current process is a child process started by multiprocessing.
    """
    return multiprocessing.current_process().name == "MainProcess"


def _get_current_dr_task_concurrency(states: Iterable[TaskInstanceState]) -> Subquery:
    """Get the dag_run IDs and how many tasks are in the provided states for each one."""
    return (
        select(TI.dag_id, TI.run_id, func.count("*").label("task_per_dr_count"))
        .where(TI.state.in_(states))
        .group_by(TI.dag_id, TI.run_id)
        .subquery()
    )


class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
    """
    SchedulerJobRunner runs for a specific time interval and schedules jobs that are ready to run.

    It figures out the latest runs for each task and sees if the dependencies
    for the next schedules are met.
    If so, it creates appropriate TaskInstances and sends run commands to the
    executor. It does this for each task in each DAG and repeats.

    :param num_runs: The number of times to run the scheduling loop. If you
        have a large number of DAG files this could complete before each file
        has been parsed. -1 for unlimited times.
    :param scheduler_idle_sleep_time: The number of seconds to wait between
        polls of running processors
    :param log: override the default Logger
    """

    job_type = "SchedulerJob"

    # For a dagrun span
    #   - key: dag_run.run_id | value: span
    #   - dagrun keys will be prefixed with 'dr:'.

Review comment (bug):
This queries the same data as `ConcurrencyMap.load()` which is still called and used for the check at lines ~680-695. Worth adding a comment explaining why we keep both? (race condition protection between query time and check time?)

Code after:
        self.dag_run_active_tasks_map.clear()
        self.task_concurrency_map.clear()
        self.task_dagrun_concurrency_map.clear()
        query = session.execute(
            select(TI.dag_id, TI.task_id, TI.run_id, func.count("*"))
            .where(TI.state.in_(EXECUTION_STATES))
            .group_by(TI.task_id, TI.run_id, TI.dag_id)
        )
        for dag_id, task_id, run_id, c in query:
            self.dag_run_active_tasks_map[dag_id, run_id] += c
            self.task_concurrency_map[(dag_id, task_id)] += c
            self.task_dagrun_concurrency_map[(dag_id, run_id, task_id)] += c


def _is_parent_process() -> bool:
    """
    Whether this is a parent process.

    Return True if the current process is the parent process.
    False if the current process is a child process started by multiprocessing.
    """
    return multiprocessing.current_process().name == "MainProcess"


def _get_current_dr_task_concurrency(states: Iterable[TaskInstanceState]) -> Subquery:
    """Get the dag_run IDs and how many tasks are in the provided states for each one."""
    return (
        select(TI.dag_id, TI.run_id, func.count("*").label("task_per_dr_count"))
        .where(TI.state.in_(states))
        .group_by(TI.dag_id, TI.run_id)
        .subquery()
    )


class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
    """
    SchedulerJobRunner runs for a specific time interval and schedules jobs that are ready to run.

    It figures out the latest runs for each task and sees if the dependencies
    for the next schedules are met.
    If so, it creates appropriate TaskInstances and sends run commands to the
    executor. It does this for each task in each DAG and repeats.

    :param num_runs: The number of times to run the scheduling loop. If you
        have a large number of DAG files this could complete before each file
        has been parsed. -1 for unlimited times.
    :param scheduler_idle_sleep_time: The number of seconds to wait between
        polls of running processors
    :param log: override the default Logger
    """