diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py index eb53d6f60db68..c3115bf6750b5 100644 --- a/airflow/cli/commands/celery_command.py +++ b/airflow/cli/commands/celery_command.py @@ -35,7 +35,7 @@ from airflow import settings from airflow.configuration import conf from airflow.utils import cli as cli_utils -from airflow.utils.cli import setup_locations, setup_logging +from airflow.utils.cli import setup_locations, setup_logging, DaemonContextWrapper from airflow.utils.providers_configuration_loader import providers_configuration_loaded from airflow.utils.serve_logs import serve_logs @@ -80,7 +80,7 @@ def flower(args): stdout.truncate(0) stderr.truncate(0) - ctx = daemon.DaemonContext( + ctx = DaemonContextWrapper( pidfile=TimeoutPIDLockFile(pidfile, -1), stdout=stdout, stderr=stderr, @@ -227,7 +227,7 @@ def worker(args): stdout_handle.truncate(0) stderr_handle.truncate(0) - daemon_context = daemon.DaemonContext( + daemon_context = DaemonContextWrapper( files_preserve=[handle], umask=int(umask, 8), stdout=stdout_handle, diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py index fd25951ad3228..40fea88112910 100644 --- a/airflow/cli/commands/scheduler_command.py +++ b/airflow/cli/commands/scheduler_command.py @@ -32,7 +32,8 @@ from airflow.jobs.job import Job, run_job from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.utils import cli as cli_utils -from airflow.utils.cli import process_subdir, setup_locations, setup_logging, sigint_handler, sigquit_handler +from airflow.utils.cli import process_subdir, setup_locations, setup_logging, sigint_handler, \ + sigquit_handler, DaemonContextWrapper from airflow.utils.providers_configuration_loader import providers_configuration_loaded from airflow.utils.scheduler_health import serve_health_check @@ -69,7 +70,7 @@ def scheduler(args): stdout_handle.truncate(0) stderr_handle.truncate(0) - ctx = daemon.DaemonContext( + ctx = DaemonContextWrapper( pidfile=TimeoutPIDLockFile(pid, -1), files_preserve=[handle], stdout=stdout_handle, diff --git a/airflow/cli/commands/triggerer_command.py b/airflow/cli/commands/triggerer_command.py index 5ddb4e23b6633..a318d17f6ddfe 100644 --- a/airflow/cli/commands/triggerer_command.py +++ b/airflow/cli/commands/triggerer_command.py @@ -31,7 +31,7 @@ from airflow.jobs.job import Job, run_job from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.utils import cli as cli_utils -from airflow.utils.cli import setup_locations, setup_logging, sigint_handler, sigquit_handler +from airflow.utils.cli import setup_locations, setup_logging, sigint_handler, sigquit_handler, DaemonContextWrapper from airflow.utils.providers_configuration_loader import providers_configuration_loaded from airflow.utils.serve_logs import serve_logs @@ -68,7 +68,7 @@ def triggerer(args): stdout_handle.truncate(0) stderr_handle.truncate(0) - daemon_context = daemon.DaemonContext( + daemon_context = DaemonContextWrapper( pidfile=TimeoutPIDLockFile(pid, -1), files_preserve=[handle], stdout=stdout_handle, @@ -79,6 +79,7 @@ def triggerer(args): triggerer_job_runner = TriggererJobRunner( job=Job(heartrate=triggerer_heartrate), capacity=args.capacity ) + run_job(job=triggerer_job_runner.job, execute_callable=triggerer_job_runner._execute) else: signal.signal(signal.SIGINT, sigint_handler) diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py index 4343c797d9244..29486e8bb4e58 100644 --- a/airflow/utils/cli.py +++ b/airflow/utils/cli.py @@ -31,6 +31,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Callable, TypeVar, cast +import daemon import re2 from sqlalchemy import select @@ -381,3 +382,14 @@ def _wrapper(*args, **kwargs): logging.disable(logging.NOTSET) return cast(T, _wrapper) + + +class DaemonContextWrapper(daemon.DaemonContext): + """Wrapper around DaemonContext to reset Stats instance in daemon process.""" + + def __enter__(self): + # in daemon context stats client needs to be reinitialized. + from airflow.stats import Stats + Stats.instance = None + + super().__enter__()