From 51b2dc6561e61ad97ac088e0d58e2b6ed317c476 Mon Sep 17 00:00:00 2001 From: Pavan Sharma Date: Thu, 5 Oct 2023 00:01:39 +0530 Subject: [PATCH 1/5] Reset stats logger instance in daemon context to reinitialize in child process --- airflow/cli/commands/celery_command.py | 6 ++++++ airflow/cli/commands/scheduler_command.py | 6 ++++++ airflow/cli/commands/triggerer_command.py | 6 ++++++ 3 files changed, 18 insertions(+) diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py index eb53d6f60db68..1199b063abf5d 100644 --- a/airflow/cli/commands/celery_command.py +++ b/airflow/cli/commands/celery_command.py @@ -87,6 +87,12 @@ def flower(args): umask=int(settings.DAEMON_UMASK, 8), ) with ctx: + + # in daemon context stats client needs to be reinitialized. + from airflow.stats import Stats + Stats.instance = None + print('Stats reset done', file=stderr) + celery_app.start(options) else: celery_app.start(options) diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py index fd25951ad3228..ba331d37bcd13 100644 --- a/airflow/cli/commands/scheduler_command.py +++ b/airflow/cli/commands/scheduler_command.py @@ -77,6 +77,12 @@ def scheduler(args): umask=int(settings.DAEMON_UMASK, 8), ) with ctx: + + # in daemon context stats client needs to be reinitialized. + from airflow.stats import Stats + Stats.instance = None + print('Stats reset done', file=stderr_handle) + _run_scheduler_job(job_runner, skip_serve_logs=args.skip_serve_logs) else: signal.signal(signal.SIGINT, sigint_handler) diff --git a/airflow/cli/commands/triggerer_command.py b/airflow/cli/commands/triggerer_command.py index 2288f1537fbb9..db7d4208c4555 100644 --- a/airflow/cli/commands/triggerer_command.py +++ b/airflow/cli/commands/triggerer_command.py @@ -77,6 +77,12 @@ def triggerer(args): umask=int(settings.DAEMON_UMASK, 8), ) with daemon_context, _serve_logs(args.skip_serve_logs): + + # in daemon context stats client needs to be reinitialized. + from airflow.stats import Stats + Stats.instance = None + print('Stats reset done', file=stderr_handle) + run_job(job=triggerer_job_runner.job, execute_callable=triggerer_job_runner._execute) else: signal.signal(signal.SIGINT, sigint_handler) From 428d54d208030db878f1d023fa11450abce13d42 Mon Sep 17 00:00:00 2001 From: Pavan Sharma Date: Mon, 16 Oct 2023 12:44:31 +0530 Subject: [PATCH 2/5] removed logs to fix test --- airflow/cli/commands/celery_command.py | 1 - airflow/cli/commands/scheduler_command.py | 1 - airflow/cli/commands/triggerer_command.py | 1 - 3 files changed, 3 deletions(-) diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py index 1199b063abf5d..1557ebced848a 100644 --- a/airflow/cli/commands/celery_command.py +++ b/airflow/cli/commands/celery_command.py @@ -91,7 +91,6 @@ def flower(args): # in daemon context stats client needs to be reinitialized. from airflow.stats import Stats Stats.instance = None - print('Stats reset done', file=stderr) celery_app.start(options) else: diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py index ba331d37bcd13..21b6929ae2234 100644 --- a/airflow/cli/commands/scheduler_command.py +++ b/airflow/cli/commands/scheduler_command.py @@ -81,7 +81,6 @@ def scheduler(args): # in daemon context stats client needs to be reinitialized. from airflow.stats import Stats Stats.instance = None - print('Stats reset done', file=stderr_handle) _run_scheduler_job(job_runner, skip_serve_logs=args.skip_serve_logs) else: diff --git a/airflow/cli/commands/triggerer_command.py b/airflow/cli/commands/triggerer_command.py index 3ac4be7c1fe7d..ba78320049d7e 100644 --- a/airflow/cli/commands/triggerer_command.py +++ b/airflow/cli/commands/triggerer_command.py @@ -80,7 +80,6 @@ def triggerer(args): # in daemon context stats client needs to be reinitialized. from airflow.stats import Stats Stats.instance = None - print('Stats reset done', file=stderr_handle) triggerer_job_runner = TriggererJobRunner( From 0936377af6b43d5885c096aea6a0a35688380dcc Mon Sep 17 00:00:00 2001 From: Pavan Sharma Date: Mon, 16 Oct 2023 19:25:17 +0530 Subject: [PATCH 3/5] added daemon contex wrapper --- airflow/cli/commands/celery_command.py | 4 ++-- airflow/cli/commands/scheduler_command.py | 10 +++------- airflow/cli/commands/triggerer_command.py | 10 ++-------- airflow/utils/cli.py | 11 +++++++++++ 4 files changed, 18 insertions(+), 17 deletions(-) diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py index 1557ebced848a..3a2113ad90c84 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, diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py index 21b6929ae2234..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, @@ -77,11 +78,6 @@ def scheduler(args): umask=int(settings.DAEMON_UMASK, 8), ) with ctx: - - # in daemon context stats client needs to be reinitialized. - from airflow.stats import Stats - Stats.instance = None - _run_scheduler_job(job_runner, skip_serve_logs=args.skip_serve_logs) else: signal.signal(signal.SIGINT, sigint_handler) diff --git a/airflow/cli/commands/triggerer_command.py b/airflow/cli/commands/triggerer_command.py index ba78320049d7e..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, @@ -76,12 +76,6 @@ def triggerer(args): umask=int(settings.DAEMON_UMASK, 8), ) with daemon_context, _serve_logs(args.skip_serve_logs): - - # in daemon context stats client needs to be reinitialized. - from airflow.stats import Stats - Stats.instance = None - - triggerer_job_runner = TriggererJobRunner( job=Job(heartrate=triggerer_heartrate), capacity=args.capacity ) diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py index 4343c797d9244..714be0093ad62 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,13 @@ def _wrapper(*args, **kwargs): logging.disable(logging.NOTSET) return cast(T, _wrapper) + + +class DaemonContextWrapper(daemon.DaemonContext): + + def __enter__(self): + # in daemon context stats client needs to be reinitialized. + from airflow.stats import Stats + Stats.instance = None + + super().__enter__() From a4196cb4bab4ea6532f6342425b35b8a932e7adf Mon Sep 17 00:00:00 2001 From: Pavan Sharma Date: Mon, 16 Oct 2023 19:28:14 +0530 Subject: [PATCH 4/5] fixed celery command --- airflow/cli/commands/celery_command.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/airflow/cli/commands/celery_command.py b/airflow/cli/commands/celery_command.py index 3a2113ad90c84..c3115bf6750b5 100644 --- a/airflow/cli/commands/celery_command.py +++ b/airflow/cli/commands/celery_command.py @@ -87,11 +87,6 @@ def flower(args): umask=int(settings.DAEMON_UMASK, 8), ) with ctx: - - # in daemon context stats client needs to be reinitialized. - from airflow.stats import Stats - Stats.instance = None - celery_app.start(options) else: celery_app.start(options) @@ -232,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, From c63d7b71630f4d5e910b8afdc6d2869fe003c24c Mon Sep 17 00:00:00 2001 From: Pavan Sharma Date: Mon, 23 Oct 2023 22:32:16 +0530 Subject: [PATCH 5/5] fix ci --- airflow/utils/cli.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py index 714be0093ad62..29486e8bb4e58 100644 --- a/airflow/utils/cli.py +++ b/airflow/utils/cli.py @@ -385,6 +385,7 @@ def _wrapper(*args, **kwargs): 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.