def _create_scheduler_job(args):
job = SchedulerJob(
subdir=process_subdir(args.subdir),
num_runs=args.num_runs,
do_pickle=args.do_pickle,
)
return job
def scheduler(args):
"""Starts Airflow Scheduler"""
print(settings.HEADER)
job = SchedulerJob(subdir=process_subdir(args.subdir),
num_runs=args.num_runs,
do_pickle=args.do_pickle)
if args.daemon:
pid, stdout, stderr, log_file = setup_locations(
"scheduler", args.pid, args.stdout, args.stderr, args.log_file)
handle = setup_logging(log_file)
stdout = open(stdout, 'w+')
stderr = open(stderr, 'w+')
ctx = daemon.DaemonContext(
pidfile=TimeoutPIDLockFile(pid, -1),
files_preserve=[handle],
stdout=stdout,
stderr=stderr,
)
with ctx:
job.run()
stdout.close()
stderr.close()
else:
signal.signal(signal.SIGINT, sigint_handler)
signal.signal(signal.SIGTERM, sigint_handler)
signal.signal(signal.SIGQUIT, sigquit_handler)
job.run()
def test_should_raise_exception_for_multiple_scheduler_on_one_host(self):
scheduler_jobs = []
with create_session() as session:
for _ in range(3):
scheduler_job = SchedulerJob()
scheduler_job.state = State.RUNNING
scheduler_job.hostname = 'HOSTNAME'
session.add(scheduler_job)
session.commit()
scheduler_job.heartbeat()
with pytest.raises(SystemExit,
match=r"Found 3 alive jobs. Expected only one."):
jobs_command.check(
self.parser.parse_args([
'jobs',
'check',
'--job-type',
'SchedulerJob',
'--limit',
'100',
]))
for scheduler_job in scheduler_jobs:
if scheduler_job.processor_agent:
scheduler_job.processor_agent.end()
def main():
try:
from airflow import conf
except ImportError:
from airflow.configuration import conf
from airflow.jobs.scheduler_job import SchedulerJob
from airflow.models import DagBag
from dbnd import dbnd_bootstrap
from dbnd._core.log.logging_utils import create_file_handler
from dbnd_airflow.executors.simple_executor import InProcessExecutor
from test_dbnd_airflow.scenarios.scheduler_perf_experiment import (
dag_folder,
dag_id,
log_scheduler,
)
dbnd_bootstrap()
conf.set("core", "unit_test_mode", "True")
logging.root.addHandler(create_file_handler(log_file=log_scheduler))
dag_bag = DagBag(dag_folder=dag_folder)
scheduler_job = SchedulerJob(
dag_ids=[dag_id],
subdir=dag_folder,
do_pickle=False,
num_runs=3,
executor=InProcessExecutor(dag_bag=dag_bag),
)
scheduler_job.run()
def scheduler(args):
"""Starts Airflow Scheduler"""
skip_serve_logs = args.skip_serve_logs
print(settings.HEADER)
job = SchedulerJob(
subdir=process_subdir(args.subdir),
num_runs=args.num_runs,
do_pickle=args.do_pickle,
)
if args.daemon:
pid, stdout, stderr, log_file = setup_locations(
"scheduler", args.pid, args.stdout, args.stderr, args.log_file)
handle = setup_logging(log_file)
with open(stdout, 'w+') as stdout_handle, open(stderr,
'w+') as stderr_handle:
ctx = daemon.DaemonContext(
pidfile=TimeoutPIDLockFile(pid, -1),
files_preserve=[handle],
stdout=stdout_handle,
stderr=stderr_handle,
)
with ctx:
sub_proc = _serve_logs(skip_serve_logs)
job.run()
else:
signal.signal(signal.SIGINT, sigint_handler)
signal.signal(signal.SIGTERM, sigint_handler)
signal.signal(signal.SIGQUIT, sigquit_handler)
sub_proc = _serve_logs(skip_serve_logs)
job.run()
if sub_proc:
sub_proc.terminate()
def run_scheduler_job(with_db_reset=False) -> None:
"""
Run the scheduler job, selectively resetting the db before creating a ScheduleJob instance
"""
from airflow.jobs.scheduler_job import SchedulerJob
if with_db_reset:
reset_db()
SchedulerJob(subdir=DAG_FOLDER, do_pickle=False, num_runs=3).run()
def test_should_ignore_not_running_jobs(self):
with create_session() as session:
for _ in range(3):
job = SchedulerJob()
job.state = State.SHUTDOWN
session.add(job)
session.commit()
# No alive jobs found.
with pytest.raises(SystemExit, match=r"No alive jobs found."):
jobs_command.check(self.parser.parse_args(['jobs', 'check']))
def test_should_report_success_for_one_working_scheduler(self):
with create_session() as session:
job = SchedulerJob()
job.state = State.RUNNING
session.add(job)
session.commit()
job.heartbeat()
with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
jobs_command.check(
self.parser.parse_args(
['jobs', 'check', '--job-type', 'SchedulerJob']))
self.assertIn("Found one alive job.", temp_stdout.getvalue())
def test_should_ignore_not_running_jobs(self):
scheduler_jobs = []
with create_session() as session:
for _ in range(3):
scheduler_job = SchedulerJob()
scheduler_job.state = State.SHUTDOWN
session.add(scheduler_job)
scheduler_jobs.append(scheduler_job)
session.commit()
# No alive jobs found.
with pytest.raises(SystemExit, match=r"No alive jobs found."):
jobs_command.check(self.parser.parse_args(['jobs', 'check']))
for scheduler_job in scheduler_jobs:
if scheduler_job.processor_agent:
scheduler_job.processor_agent.end()
def test_should_report_success_for_ha_schedulers(self):
with create_session() as session:
for _ in range(3):
job = SchedulerJob()
job.state = State.RUNNING
session.add(job)
session.commit()
job.heartbeat()
with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
jobs_command.check(
self.parser.parse_args([
'jobs', 'check', '--job-type', 'SchedulerJob', '--limit',
'100', '--allow-multiple'
]))
self.assertIn("Found 3 alive jobs.", temp_stdout.getvalue())
def scheduler(args):
"""Starts Airflow Scheduler"""
print(settings.HEADER)
job = SchedulerJob(
subdir=process_subdir(args.subdir),
num_runs=args.num_runs,
do_pickle=args.do_pickle,
)
scheduler_name = SchedulerFactory.get_scheduler_name()
if scheduler_name == SchedulerFactory.DEFAULT_SCHEDULER:
pass
elif scheduler_name == SchedulerFactory.EVENT_BASED_SCHEDULER:
job = EventBasedSchedulerJob(dag_directory=process_subdir(args.subdir),
server_uri=args.server_uri)
else:
scheduler_class = SchedulerFactory.get_default_scheduler()
job = scheduler_class()
if args.daemon:
pid, stdout, stderr, log_file = setup_locations(
"scheduler", args.pid, args.stdout, args.stderr, args.log_file)
handle = setup_logging(log_file)
stdout = open(stdout, 'w+')
stderr = open(stderr, 'w+')
ctx = daemon.DaemonContext(
pidfile=TimeoutPIDLockFile(pid, -1),
files_preserve=[handle],
stdout=stdout,
stderr=stderr,
)
with ctx:
job.run()
stdout.close()
stderr.close()
else:
signal.signal(signal.SIGINT, sigint_handler)
signal.signal(signal.SIGTERM, sigint_handler)
signal.signal(signal.SIGQUIT, sigquit_handler)
job.run()
def main(num_runs, repeat, pre_create_dag_runs, executor_class, dag_ids):
"""
This script can be used to measure the total "scheduler overhead" of Airflow.
By overhead we mean if the tasks executed instantly as soon as they are
executed (i.e. they do nothing) how quickly could we schedule them.
It will monitor the task completion of the Mock/stub executor (no actual
tasks are run) and after the required number of dag runs for all the
specified dags have completed all their tasks, it will cleanly shut down
the scheduler.
The dags you run with need to have an early enough start_date to create the
desired number of runs.
Care should be taken that other limits (DAG concurrency, pool size etc) are
not the bottleneck. This script doesn't help you in that regard.
It is recommended to repeat the test at least 3 times (`--repeat=3`, the
default) so that you can get somewhat-accurate variance on the reported
timing numbers, but this can be disabled for longer runs if needed.
"""
# Turn on unit test mode so that we don't do any sleep() in the scheduler
# loop - not needed on master, but this script can run against older
# releases too!
os.environ['AIRFLOW__CORE__UNIT_TEST_MODE'] = 'True'
os.environ['AIRFLOW__CORE__DAG_CONCURRENCY'] = '500'
# Set this so that dags can dynamically configure their end_date
os.environ['AIRFLOW_BENCHMARK_MAX_DAG_RUNS'] = str(num_runs)
os.environ['PERF_MAX_RUNS'] = str(num_runs)
if pre_create_dag_runs:
os.environ['AIRFLOW__SCHEDULER__USE_JOB_SCHEDULE'] = 'False'
from airflow.jobs.scheduler_job import SchedulerJob
from airflow.models.dagbag import DagBag
from airflow.utils import db
dagbag = DagBag()
dags = []
with db.create_session() as session:
pause_all_dags(session)
for dag_id in dag_ids:
dag = dagbag.get_dag(dag_id)
dag.sync_to_db(session=session)
dags.append(dag)
reset_dag(dag, session)
next_run_date = dag.normalize_schedule(dag.start_date
or min(t.start_date
for t in dag.tasks))
for _ in range(num_runs - 1):
next_run_date = dag.following_schedule(next_run_date)
end_date = dag.end_date or dag.default_args.get('end_date')
if end_date != next_run_date:
message = (
f"DAG {dag_id} has incorrect end_date ({end_date}) for number of runs! "
f"It should be "
f" {next_run_date}")
sys.exit(message)
if pre_create_dag_runs:
create_dag_runs(dag, num_runs, session)
ShortCircuitExecutor = get_executor_under_test(executor_class)
executor = ShortCircuitExecutor(dag_ids_to_watch=dag_ids,
num_runs=num_runs)
scheduler_job = SchedulerJob(dag_ids=dag_ids,
do_pickle=False,
executor=executor)
executor.scheduler_job = scheduler_job
total_tasks = sum(len(dag.tasks) for dag in dags)
if 'PYSPY' in os.environ:
pid = str(os.getpid())
filename = os.environ.get('PYSPY_O', 'flame-' + pid + '.html')
os.spawnlp(os.P_NOWAIT, 'sudo', 'sudo', 'py-spy', 'record', '-o',
filename, '-p', pid, '--idle')
times = []
# Need a lambda to refer to the _latest_ value for scheduler_job, not just
# the initial one
code_to_test = lambda: scheduler_job.run() # pylint: disable=unnecessary-lambda
for count in range(repeat):
gc.disable()
start = time.perf_counter()
code_to_test()
times.append(time.perf_counter() - start)
gc.enable()
print("Run %d time: %.5f" % (count + 1, times[-1]))
if count + 1 != repeat:
with db.create_session() as session:
for dag in dags:
reset_dag(dag, session)
executor.reset(dag_ids)
scheduler_job = SchedulerJob(dag_ids=dag_ids,
do_pickle=False,
executor=executor)
executor.scheduler_job = scheduler_job
print()
print()
msg = "Time for %d dag runs of %d dags with %d total tasks: %.4fs"
if len(times) > 1:
print((msg + " (±%.3fs)") %
(num_runs, len(dags), total_tasks, statistics.mean(times),
statistics.stdev(times)))
else:
print(msg % (num_runs, len(dags), total_tasks, times[0]))
print()
print()
def run_scheduler_job(with_db_reset=False) -> None:
from airflow.jobs.scheduler_job import SchedulerJob
if with_db_reset:
reset_db()
SchedulerJob(subdir=DAG_FOLDER, do_pickle=False, num_runs=3).run()
