def test_dag_run_event_manager(self): mailbox = Mailbox() event_manager = DagRunEventManager(mailbox) event_manager.start() self.create_task_state(dag_run=self._dag_run, task_id='operator_toggle_handler') event = BaseEvent("test_event", "test_event", namespace="default") event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) handle_event = mailbox.get_message() message = EventHandleResult.from_event(handle_event) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.START) handle_event = mailbox.get_message() message = EventHandleResult.from_event(handle_event) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.STOP) time.sleep(2) event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) handle_event = mailbox.get_message() message = EventHandleResult.from_event(handle_event) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.START) event_manager.end()
def test_dag_run_event_manager_resubmit_if_exit_with_nonempty_queue(self): mailbox = Mailbox() event_manager = DagRunEventManager(mailbox, max_num_event=1) event_manager.start() self.create_task_state(dag_run=self._dag_run, task_id='operator_toggle_handler') event = BaseEvent("test_event", "test_event", namespace="default") event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) assert mailbox.get_message() is not None assert mailbox.get_message_with_timeout(5) is not None
class EventBasedSchedulerJob(BaseJob): """ 1. todo self heartbeat 2. todo check other scheduler failed 3. todo timeout dagrun """ __mapper_args__ = {'polymorphic_identity': 'EventBasedSchedulerJob'} def __init__(self, dag_directory, server_uri=None, max_runs=-1, refresh_dag_dir_interval=conf.getint( 'scheduler', 'refresh_dag_dir_interval', fallback=30), *args, **kwargs): super().__init__(*args, **kwargs) self.mailbox: Mailbox = Mailbox() self.dag_trigger: DagTrigger = DagTrigger( dag_directory=dag_directory, max_runs=max_runs, dag_ids=None, pickle_dags=False, mailbox=self.mailbox, refresh_dag_dir_interval=refresh_dag_dir_interval, notification_service_uri=server_uri) self.task_event_manager = DagRunEventManager(self.mailbox) self.executor.set_mailbox(self.mailbox) self.notification_client: NotificationClient = NotificationClient( server_uri=server_uri, default_namespace=SCHEDULER_NAMESPACE) self.scheduler: EventBasedScheduler = EventBasedScheduler( self.id, self.mailbox, self.task_event_manager, self.executor, self.notification_client) self.last_scheduling_id = self._last_scheduler_job_id() @staticmethod def _last_scheduler_job_id(): last_run = EventBasedSchedulerJob.most_recent_job() if not last_run: return None else: return last_run.id def _execute(self): # faulthandler.enable() self.log.info("Starting the scheduler Job") # DAGs can be pickled for easier remote execution by some executors # pickle_dags = self.do_pickle and self.executor_class not in UNPICKLEABLE_EXECUTORS try: self.mailbox.set_scheduling_job_id(self.id) self.scheduler.id = self.id self._start_listen_events() self.dag_trigger.start() self.task_event_manager.start() self.executor.job_id = self.id self.executor.start() self.register_signals() # Start after resetting orphaned tasks to avoid stressing out DB. execute_start_time = timezone.utcnow() self.scheduler.submit_sync_thread() self.scheduler.recover(self.last_scheduling_id) self.scheduler.schedule() self.executor.end() self.dag_trigger.end() self.task_event_manager.end() self._stop_listen_events() settings.Session.remove() # type: ignore except Exception as e: # pylint: disable=broad-except self.log.exception("Exception when executing scheduler, %s", e) finally: self.log.info("Exited execute loop") def _start_listen_events(self): watcher = SchedulerEventWatcher(self.mailbox) self.notification_client.start_listen_events(watcher=watcher, start_time=int( time.time() * 1000), version=None) def _stop_listen_events(self): self.notification_client.stop_listen_events() def register_signals(self) -> None: """Register signals that stop child processes""" signal.signal(signal.SIGINT, self._exit_gracefully) signal.signal(signal.SIGTERM, self._exit_gracefully) signal.signal(signal.SIGUSR2, self._debug_dump) def _exit_gracefully(self, signum, frame) -> None: # pylint: disable=unused-argument """Helper method to clean up processor_agent to avoid leaving orphan processes.""" self.log.info("Exiting gracefully upon receiving signal %s", signum) sys.exit(os.EX_OK) def _debug_dump(self, signum, frame): # pylint: disable=unused-argument try: sig_name = signal.Signals(signum).name # pylint: disable=no-member except Exception: # pylint: disable=broad-except sig_name = str(signum) self.log.info("%s\n%s received, printing debug\n%s", "-" * 80, sig_name, "-" * 80) self.executor.debug_dump() self.log.info("-" * 80) def is_alive(self, grace_multiplier: Optional[float] = None) -> bool: """ Is this SchedulerJob alive? We define alive as in a state of running and a heartbeat within the threshold defined in the ``scheduler_health_check_threshold`` config setting. ``grace_multiplier`` is accepted for compatibility with the parent class. :rtype: boolean """ if grace_multiplier is not None: # Accept the same behaviour as superclass return super().is_alive(grace_multiplier=grace_multiplier) scheduler_health_check_threshold: int = conf.getint( 'scheduler', 'scheduler_health_check_threshold') return (self.state == State.RUNNING and (timezone.utcnow() - self.latest_heartbeat).total_seconds() < scheduler_health_check_threshold)