def test_set_file_paths_when_processor_file_path_is_in_new_file_paths(self): manager = DagFileProcessorManager(dag_directory='directory', file_paths=['abc.txt'], parallelism=1, process_file_interval=1, max_runs=1, processor_factory=MagicMock().return_value) mock_processor = MagicMock() mock_processor.stop.side_effect = AttributeError( 'DagFileProcessor object has no attribute stop') mock_processor.terminate.side_effect = None manager._processors['abc.txt'] = mock_processor manager.set_file_paths(['abc.txt']) self.assertDictEqual(manager._processors, {'abc.txt': mock_processor})
def test_kill_timed_out_processors_kill(self, mock_kill, mock_pid): mock_pid.return_value = 1234 manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta(seconds=5), signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True, ) processor = DagFileProcessorProcess('abc.txt', False, [], []) processor._start_time = timezone.make_aware(datetime.min) manager._processors = {'abc.txt': processor} manager._kill_timed_out_processors() mock_kill.assert_called_once_with()
def test_find_zombies(self): manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), async_mode=True) dagbag = DagBag(TEST_DAG_FOLDER) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') dag.sync_to_db() task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) local_job = LJ(ti) local_job.state = State.SHUTDOWN local_job.id = 1 ti.job_id = local_job.id session.add(local_job) session.add(ti) session.commit() manager._last_zombie_query_time = timezone.utcnow() - timedelta( seconds=manager._zombie_threshold_secs + 1) manager._find_zombies() # pylint: disable=no-value-for-parameter requests = manager._callback_to_execute[dag.full_filepath] self.assertEqual(1, len(requests)) self.assertEqual(requests[0].full_filepath, dag.full_filepath) self.assertEqual(requests[0].msg, "Detected as zombie") self.assertIsInstance(requests[0].simple_task_instance, SimpleTaskInstance) self.assertEqual(ti.dag_id, requests[0].simple_task_instance.dag_id) self.assertEqual(ti.task_id, requests[0].simple_task_instance.task_id) self.assertEqual(ti.execution_date, requests[0].simple_task_instance.execution_date) session.query(TI).delete() session.query(LJ).delete()
def test_dag_with_system_exit(self): """ Test to check that a DAG with a system.exit() doesn't break the scheduler. """ # We need to _actually_ parse the files here to test the behaviour. # Right now the parsing code lives in SchedulerJob, even though it's # called via utils.dag_processing. from airflow.jobs.scheduler_job import SchedulerJob dag_id = 'exit_test_dag' dag_directory = TEST_DAG_FOLDER.parent / 'dags_with_system_exit' # Delete the one valid DAG/SerializedDAG, and check that it gets re-created clear_db_dags() clear_db_serialized_dags() child_pipe, parent_pipe = multiprocessing.Pipe() manager = DagFileProcessorManager( dag_directory=dag_directory, dag_ids=[], max_runs=1, processor_factory=SchedulerJob._create_dag_file_processor, processor_timeout=timedelta(seconds=5), signal_conn=child_pipe, pickle_dags=False, async_mode=True, ) manager._run_parsing_loop() result = None while parent_pipe.poll(timeout=None): result = parent_pipe.recv() if isinstance(result, DagParsingStat) and result.done: break # Three files in folder should be processed assert sum(stat.run_count for stat in manager._file_stats.values()) == 3 with create_session() as session: assert session.query(DagModel).get(dag_id) is not None
def test_set_file_paths_when_processor_file_path_is_in_new_file_paths( self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], parallelism=1, process_file_interval=1, max_runs=1, processor_factory=MagicMock().return_value) mock_processor = MagicMock() mock_processor.stop.side_effect = AttributeError( 'DagFileProcessor object has no attribute stop') mock_processor.terminate.side_effect = None manager._processors['abc.txt'] = mock_processor manager.set_file_paths(['abc.txt']) self.assertDictEqual(manager._processors, {'abc.txt': mock_processor})
def test_kill_timed_out_processors_no_kill(self, mock_dag_file_processor, mock_pid): mock_pid.return_value = 1234 manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta(seconds=5), signal_conn=MagicMock(), stat_queue=MagicMock(), result_queue=MagicMock, async_mode=True) processor = DagFileProcessor('abc.txt', False, []) processor._start_time = timezone.make_aware(datetime.max) manager._processors = {'abc.txt': processor} manager._kill_timed_out_processors() mock_dag_file_processor.kill.assert_not_called()
def test_set_file_paths_when_processor_file_path_is_in_new_file_paths( self): manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), async_mode=True) mock_processor = MagicMock() mock_processor.stop.side_effect = AttributeError( 'DagFileProcessor object has no attribute stop') mock_processor.terminate.side_effect = None manager._processors['abc.txt'] = mock_processor manager.set_file_paths(['abc.txt']) self.assertDictEqual(manager._processors, {'abc.txt': mock_processor})
def test_set_file_paths_when_processor_file_path_is_in_new_file_paths(self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, signal_conn=MagicMock(), stat_queue=MagicMock(), result_queue=MagicMock, async_mode=True) mock_processor = MagicMock() mock_processor.stop.side_effect = AttributeError( 'DagFileProcessor object has no attribute stop') mock_processor.terminate.side_effect = None manager._processors['abc.txt'] = mock_processor manager.set_file_paths(['abc.txt']) self.assertDictEqual(manager._processors, {'abc.txt': mock_processor})
def test_set_file_paths_when_processor_file_path_not_in_new_file_paths(self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, signal_conn=MagicMock(), stat_queue=MagicMock(), result_queue=MagicMock, async_mode=True) mock_processor = MagicMock() mock_processor.stop.side_effect = AttributeError( 'DagFileProcessor object has no attribute stop') mock_processor.terminate.side_effect = None manager._processors['missing_file.txt'] = mock_processor manager.set_file_paths(['abc.txt']) self.assertDictEqual(manager._processors, {})
def test_set_file_paths_when_processor_file_path_not_in_new_file_paths(self): manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True, ) mock_processor = MagicMock() mock_processor.stop.side_effect = AttributeError('DagFileProcessor object has no attribute stop') mock_processor.terminate.side_effect = None manager._processors['missing_file.txt'] = mock_processor manager._file_stats['missing_file.txt'] = DagFileStat(0, 0, None, None, 0) manager.set_file_paths(['abc.txt']) assert manager._processors == {}
def test_find_zombies(self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True) dagbag = DagBag(TEST_DAG_FOLDER) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() manager._last_zombie_query_time = timezone.utcnow() - timedelta( seconds=manager._zombie_threshold_secs + 1) manager._find_zombies() zombies = manager._zombies self.assertEqual(1, len(zombies)) self.assertIsInstance(zombies[0], SimpleTaskInstance) self.assertEqual(ti.dag_id, zombies[0].dag_id) self.assertEqual(ti.task_id, zombies[0].task_id) self.assertEqual(ti.execution_date, zombies[0].execution_date) session.query(TI).delete() session.query(LJ).delete()
def test_file_paths_in_queue_sorted_by_modified_time( self, mock_getmtime, mock_isfile, mock_find_path, mock_might_contain_dag, mock_zipfile): """Test files are sorted by modified time""" paths_with_mtime = { "file_3.py": 3.0, "file_2.py": 2.0, "file_4.py": 5.0, "file_1.py": 4.0 } dag_files = list(paths_with_mtime.keys()) mock_getmtime.side_effect = list(paths_with_mtime.values()) mock_find_path.return_value = dag_files manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True, ) manager.set_file_paths(dag_files) assert manager._file_path_queue == [] manager.prepare_file_path_queue() assert manager._file_path_queue == [ 'file_4.py', 'file_1.py', 'file_3.py', 'file_2.py' ]
def test_file_paths_in_queue_sorted_alphabetically(self, mock_isfile, mock_find_path, mock_might_contain_dag, mock_zipfile): """Test dag files are sorted alphabetically""" dag_files = ["file_3.py", "file_2.py", "file_4.py", "file_1.py"] mock_find_path.return_value = dag_files manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True, ) manager.set_file_paths(dag_files) assert manager._file_path_queue == [] manager.prepare_file_path_queue() assert manager._file_path_queue == [ 'file_1.py', 'file_2.py', 'file_3.py', 'file_4.py' ]
def test_find_zombies(self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, signal_conn=MagicMock(), stat_queue=MagicMock(), result_queue=MagicMock, async_mode=True) dagbag = DagBag(TEST_DAG_FOLDER) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('example_branch_operator') task = dag.get_task(task_id='run_this_first') ti = TI(task, DEFAULT_DATE, State.RUNNING) lj = LJ(ti) lj.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() manager._last_zombie_query_time = timezone.utcnow() - timedelta( seconds=manager._zombie_threshold_secs + 1) zombies = manager._find_zombies() self.assertEqual(1, len(zombies)) self.assertIsInstance(zombies[0], SimpleTaskInstance) self.assertEqual(ti.dag_id, zombies[0].dag_id) self.assertEqual(ti.task_id, zombies[0].task_id) self.assertEqual(ti.execution_date, zombies[0].execution_date) session.query(TI).delete() session.query(LJ).delete()
def test_cleanup_stale_dags_no_serialization(self, sdm_mock, dag_mock): manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=MagicMock().return_value, processor_timeout=timedelta(seconds=50), signal_conn=MagicMock(), async_mode=True) manager.last_dag_cleanup_time = DEFAULT_DATE - timezone.dt.timedelta(seconds=301) manager._file_process_interval = 30 manager._min_serialized_dag_update_interval = 30 expected_min_last_seen = DEFAULT_DATE - timezone.dt.timedelta(seconds=(50 + 30 + 30)) manager._cleanup_stale_dags() dag_mock.deactivate_stale_dags.assert_called_with(expected_min_last_seen) sdm_mock.remove_stale_dags.assert_called_with(expected_min_last_seen)
def test_max_runs_when_no_files(self): child_pipe, parent_pipe = multiprocessing.Pipe() with TemporaryDirectory(prefix="empty-airflow-dags-") as dags_folder: async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn') manager = DagFileProcessorManager( dag_directory=dags_folder, max_runs=1, processor_factory=FakeDagFileProcessorRunner._fake_dag_processor_factory, processor_timeout=timedelta.max, signal_conn=child_pipe, dag_ids=[], pickle_dags=False, async_mode=async_mode) self.run_processor_manager_one_loop(manager, parent_pipe) child_pipe.close() parent_pipe.close()
def test_start_new_processes_with_same_filepath(self): """ Test that when a processor already exist with a filepath, a new processor won't be created with that filepath. The filepath will just be removed from the list. """ processor_factory_mock = MagicMock() manager = DagFileProcessorManager( dag_directory='directory', max_runs=1, processor_factory=processor_factory_mock, processor_timeout=timedelta.max, signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True, ) file_1 = 'file_1.py' file_2 = 'file_2.py' file_3 = 'file_3.py' manager._file_path_queue = [file_1, file_2, file_3] # Mock that only one processor exists. This processor runs with 'file_1' manager._processors[file_1] = MagicMock() # Start New Processes manager.start_new_processes() # Because of the config: '[scheduler] parsing_processes = 2' # verify that only one extra process is created # and since a processor with 'file_1' already exists, # even though it is first in '_file_path_queue' # a new processor is created with 'file_2' and not 'file_1'. processor_factory_mock.assert_called_once_with('file_2.py', [], [], False) assert file_1 in manager._processors.keys() assert file_2 in manager._processors.keys() assert [file_3] == manager._file_path_queue
def test_handle_failure_callback_with_zombies_are_correctly_passed_to_dag_file_processor(self): """ Check that the same set of failure callback with zombies are passed to the dag file processors until the next zombie detection logic is invoked. """ test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py') with conf_vars({('scheduler', 'parsing_processes'): '1', ('core', 'load_examples'): 'False'}): dagbag = DagBag(test_dag_path, read_dags_from_db=False) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('test_example_bash_operator') dag.sync_to_db() task = dag.get_task(task_id='run_this_last') ti = TI(task, DEFAULT_DATE, State.RUNNING) local_job = LJ(ti) local_job.state = State.SHUTDOWN session.add(local_job) session.commit() # TODO: If there was an actual Relationship between TI and Job # we wouldn't need this extra commit session.add(ti) ti.job_id = local_job.id session.commit() expected_failure_callback_requests = [ TaskCallbackRequest( full_filepath=dag.full_filepath, simple_task_instance=SimpleTaskInstance(ti), msg="Message", ) ] test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py') child_pipe, parent_pipe = multiprocessing.Pipe() async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn') fake_processors = [] def fake_processor_factory(*args, **kwargs): nonlocal fake_processors processor = FakeDagFileProcessorRunner._fake_dag_processor_factory(*args, **kwargs) fake_processors.append(processor) return processor manager = DagFileProcessorManager( dag_directory=test_dag_path, max_runs=1, processor_factory=fake_processor_factory, processor_timeout=timedelta.max, signal_conn=child_pipe, dag_ids=[], pickle_dags=False, async_mode=async_mode, ) self.run_processor_manager_one_loop(manager, parent_pipe) if async_mode: # Once for initial parse, and then again for the add_callback_to_queue assert len(fake_processors) == 2 assert fake_processors[0]._file_path == test_dag_path assert fake_processors[0]._callback_requests == [] else: assert len(fake_processors) == 1 assert fake_processors[-1]._file_path == test_dag_path callback_requests = fake_processors[-1]._callback_requests assert {zombie.simple_task_instance.key for zombie in expected_failure_callback_requests} == { result.simple_task_instance.key for result in callback_requests } child_pipe.close() parent_pipe.close()
def test_recently_modified_file_is_parsed_with_mtime_mode( self, mock_getmtime, mock_isfile, mock_find_path, mock_might_contain_dag, mock_zipfile ): """ Test recently updated files are processed even if min_file_process_interval is not reached """ freezed_base_time = timezone.datetime(2020, 1, 5, 0, 0, 0) initial_file_1_mtime = (freezed_base_time - timedelta(minutes=5)).timestamp() dag_files = ["file_1.py"] mock_getmtime.side_effect = [initial_file_1_mtime] mock_find_path.return_value = dag_files manager = DagFileProcessorManager( dag_directory='directory', max_runs=3, processor_factory=MagicMock().return_value, processor_timeout=timedelta.max, signal_conn=MagicMock(), dag_ids=[], pickle_dags=False, async_mode=True, ) # let's say the DAG was just parsed 2 seconds before the Freezed time last_finish_time = freezed_base_time - timedelta(seconds=10) manager._file_stats = { "file_1.py": DagFileStat(1, 0, last_finish_time, 1.0, 1), } with freeze_time(freezed_base_time): manager.set_file_paths(dag_files) assert manager._file_path_queue == [] # File Path Queue will be empty as the "modified time" < "last finish time" manager.prepare_file_path_queue() assert manager._file_path_queue == [] # Simulate the DAG modification by using modified_time which is greater # than the last_parse_time but still less than now - min_file_process_interval file_1_new_mtime = freezed_base_time - timedelta(seconds=5) file_1_new_mtime_ts = file_1_new_mtime.timestamp() with freeze_time(freezed_base_time): manager.set_file_paths(dag_files) assert manager._file_path_queue == [] # File Path Queue will be empty as the "modified time" < "last finish time" mock_getmtime.side_effect = [file_1_new_mtime_ts] manager.prepare_file_path_queue() # Check that file is added to the queue even though file was just recently passed assert manager._file_path_queue == ["file_1.py"] assert last_finish_time < file_1_new_mtime assert ( manager._file_process_interval > (freezed_base_time - manager.get_last_finish_time("file_1.py")).total_seconds() )
def test_handle_failure_callback_with_zombies_are_correctly_passed_to_dag_file_processor( self): """ Check that the same set of failure callback with zombies are passed to the dag file processors until the next zombie detection logic is invoked. """ test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py') with conf_vars({ ('scheduler', 'max_threads'): '1', ('core', 'load_examples'): 'False' }): dagbag = DagBag(test_dag_path) with create_session() as session: session.query(LJ).delete() dag = dagbag.get_dag('test_example_bash_operator') dag.sync_to_db() task = dag.get_task(task_id='run_this_last') ti = TI(task, DEFAULT_DATE, State.RUNNING) local_job = LJ(ti) local_job.state = State.SHUTDOWN session.add(local_job) session.commit() # TODO: If there was an actual Relationshop between TI and Job # we wouldn't need this extra commit session.add(ti) ti.job_id = local_job.id session.commit() fake_failure_callback_requests = [ FailureCallbackRequest( full_filepath=dag.full_filepath, simple_task_instance=SimpleTaskInstance(ti), msg="Message") ] test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py') child_pipe, parent_pipe = multiprocessing.Pipe() async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn') manager = DagFileProcessorManager( dag_directory=test_dag_path, max_runs=1, processor_factory=FakeDagFileProcessorRunner. _fake_dag_processor_factory, processor_timeout=timedelta.max, signal_conn=child_pipe, dag_ids=[], pickle_dags=False, async_mode=async_mode) parsing_result = self.run_processor_manager_one_loop( manager, parent_pipe) self.assertEqual(len(fake_failure_callback_requests), len(parsing_result)) self.assertEqual( set(zombie.simple_task_instance.key for zombie in fake_failure_callback_requests), set(result.simple_task_instance.key for result in parsing_result)) child_pipe.close() parent_pipe.close()
def test_pipe_full_deadlock(self): dag_filepath = TEST_DAG_FOLDER / "test_scheduler_dags.py" child_pipe, parent_pipe = multiprocessing.Pipe() # Shrink the buffers to exacerbate the problem! for fd in (parent_pipe.fileno(),): sock = socket.socket(fileno=fd) sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024) sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 1024) sock.detach() exit_event = threading.Event() # To test this behaviour we need something that continually fills the # parent pipe's buffer (and keeps it full). def keep_pipe_full(pipe, exit_event): n = 0 while True: if exit_event.is_set(): break req = CallbackRequest(str(dag_filepath)) try: logging.debug("Sending CallbackRequests %d", n + 1) pipe.send(req) except TypeError: # This is actually the error you get when the parent pipe # is closed! Nicely handled, eh? break except OSError: break n += 1 logging.debug(" Sent %d CallbackRequests", n) thread = threading.Thread(target=keep_pipe_full, args=(parent_pipe, exit_event)) fake_processors = [] def fake_processor_factory(*args, **kwargs): nonlocal fake_processors processor = FakeDagFileProcessorRunner._fake_dag_processor_factory(*args, **kwargs) fake_processors.append(processor) return processor manager = DagFileProcessorManager( dag_directory=dag_filepath, dag_ids=[], # A reasonable large number to ensure that we trigger the deadlock max_runs=100, processor_factory=fake_processor_factory, processor_timeout=timedelta(seconds=5), signal_conn=child_pipe, pickle_dags=False, async_mode=True, ) try: thread.start() # If this completes without hanging, then the test is good! manager._run_parsing_loop() exit_event.set() finally: logging.info("Closing pipes") parent_pipe.close() child_pipe.close() thread.join(timeout=1.0)