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()
def main(num_runs, repeat, pre_create_dag_runs, executor_class, dag_ids): # pylint: disable=too-many-locals """ 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 fo 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()
class TestCliConfigList(unittest.TestCase): @classmethod def setUpClass(cls): cls.parser = cli_parser.get_parser() def setUp(self) -> None: clear_db_jobs() self.scheduler_job = None def tearDown(self) -> None: if self.scheduler_job and self.scheduler_job.processor_agent: self.scheduler_job.processor_agent.end() clear_db_jobs() def test_should_report_success_for_one_working_scheduler(self): with create_session() as session: self.scheduler_job = SchedulerJob() self.scheduler_job.state = State.RUNNING session.add(self.scheduler_job) session.commit() self.scheduler_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_report_success_for_one_working_scheduler_with_hostname( self): with create_session() as session: self.scheduler_job = SchedulerJob() self.scheduler_job.state = State.RUNNING self.scheduler_job.hostname = 'HOSTNAME' session.add(self.scheduler_job) session.commit() self.scheduler_job.heartbeat() with contextlib.redirect_stdout(io.StringIO()) as temp_stdout: jobs_command.check( self.parser.parse_args([ 'jobs', 'check', '--job-type', 'SchedulerJob', '--hostname', 'HOSTNAME' ])) self.assertIn("Found one alive job.", temp_stdout.getvalue()) def test_should_report_success_for_ha_schedulers(self): scheduler_jobs = [] with create_session() as session: for _ in range(3): scheduler_job = SchedulerJob() scheduler_job.state = State.RUNNING session.add(scheduler_job) scheduler_jobs.append(scheduler_job) session.commit() scheduler_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()) for scheduler_job in scheduler_jobs: if scheduler_job.processor_agent: scheduler_job.processor_agent.end() 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_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 test_should_raise_exception_for_allow_multiple_and_limit_1(self): with pytest.raises( SystemExit, match= r"To use option --allow-multiple, you must set the limit to a value greater than 1.", ): jobs_command.check( self.parser.parse_args(['jobs', 'check', '--allow-multiple']))