Repository: flink-extended/dl-on-flink (language: python)

Issue:
[Flink AI Flow] EventBasedScheduler doesn't heartbeat
<!--

*Thank you very much for contributing to flink-ai-extended. To help the community review your issue or contribution in the best possible way，please take a few minutes to fulfill following items.*

-->

### Describe the bug
EventBasedScheduler doesn't heartbeat, so that WebUI ```displays scheduler not running```
![image](https://user-images.githubusercontent.com/21049075/119600028-b7f96700-be18-11eb-876b-1a4b0ece2294.png)


### Your environment
#### Operating system
MacOS

#### Database
<!-- The database and version you use -->
mysql8.0

#### Python version
<!-- The python version you use -->
3.7.9

### To Reproduce
Steps to reproduce the behavior (if you can):
1. Run start-aiflow.sh
2. Open WebUI

### Expected behavior
<!-- A clear and concise description of what you expected to happen -->
1. No ```scheduler not running ``` in WebUI
2. Table colume```job.last_heartbeat``` can update every several seconds.

### Actual behavior
<!-- A clear and concise description of what you see instead -->
1. Run start-aiflow.sh
2. Open WebUI, you can see scheduler not running
3. Query database table ```job.last_heartbeat```, the SchedulerJob you just started doesn't heartbeat at all.

### Screenshots
<!-- If applicable, add screenshots to help explain your problem -->

### Additional context
<!-- Please include any error logs -->
#### proposal
1. Add a configuration ```scheduler = EventBasedSchedulerJob``` with default value ```SchedulerJob``` in airflow.cfg
2. Use ```airflow scheduler``` command to start default scheduler/event_based scheduler
3. Move schedule loop from EventBasedScheduler to EventBasedSchedulerJob, send heartbeat during each loop. 



Discussion / hints:

Fix patch:
diff --git a/flink-ai-flow/ai_flow/bin/start-aiflow.sh b/flink-ai-flow/ai_flow/bin/start-aiflow.sh
index 7dc9936c1..bb2260795 100755
--- a/flink-ai-flow/ai_flow/bin/start-aiflow.sh
+++ b/flink-ai-flow/ai_flow/bin/start-aiflow.sh
@@ -43,6 +43,7 @@ if [[ ! -f "${AIRFLOW_HOME}/airflow.cfg" ]] ; then
         gsub(\"executor = SequentialExecutor\", \"executor = LocalExecutor\"); \
         gsub(\"dags_are_paused_at_creation = True\", \"dags_are_paused_at_creation = False\"); \
         gsub(\"# mp_start_method =\", \"mp_start_method = forkserver\"); \
+        gsub(\"# scheduler =\", \"scheduler = EventBasedSchedulerJob\"); \
         gsub(\"execute_tasks_new_python_interpreter = False\", \"execute_tasks_new_python_interpreter = True\"); \
         gsub(\"min_serialized_dag_update_interval = 30\", \"min_serialized_dag_update_interval = 0\"); \
         print \$0}" airflow.cfg.tmpl > airflow.cfg
@@ -80,7 +81,7 @@ start_notification_service.py --database-conn=${MYSQL_CONN} > ${AIFLOW_LOG_DIR}/
 echo $! > ${AIFLOW_PID_DIR}/notification_service.pid
 
 # start airflow scheduler and web server
-airflow event_scheduler --subdir=${AIRFLOW_DEPLOY_PATH} > ${AIFLOW_LOG_DIR}/scheduler.log 2>&1 &
+airflow scheduler --subdir=${AIRFLOW_DEPLOY_PATH} > ${AIFLOW_LOG_DIR}/scheduler.log 2>&1 &
 echo $! > ${AIFLOW_PID_DIR}/scheduler.pid
 airflow webserver -p 8080 > ${AIFLOW_LOG_DIR}/web.log 2>&1 &
 echo $! > ${AIFLOW_PID_DIR}/web.pid
diff --git a/flink-ai-flow/lib/airflow/airflow/api_connexion/endpoints/health_endpoint.py b/flink-ai-flow/lib/airflow/airflow/api_connexion/endpoints/health_endpoint.py
index ea0d62d57..265959ebd 100644
--- a/flink-ai-flow/lib/airflow/airflow/api_connexion/endpoints/health_endpoint.py
+++ b/flink-ai-flow/lib/airflow/airflow/api_connexion/endpoints/health_endpoint.py
@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 from airflow.api_connexion.schemas.health_schema import health_schema
-from airflow.jobs.scheduler_job import SchedulerJob
+from airflow.contrib.jobs.scheduler_factory import SchedulerFactory
 
 HEALTHY = "healthy"
 UNHEALTHY = "unhealthy"
@@ -27,7 +27,8 @@ def get_health():
     latest_scheduler_heartbeat = None
     scheduler_status = UNHEALTHY
     try:
-        scheduler_job = SchedulerJob.most_recent_job()
+        scheduler_class = SchedulerFactory.get_default_scheduler()
+        scheduler_job = scheduler_class.most_recent_job()
 
         if scheduler_job:
             latest_scheduler_heartbeat = scheduler_job.latest_heartbeat.isoformat()
diff --git a/flink-ai-flow/lib/airflow/airflow/cli/cli_parser.py b/flink-ai-flow/lib/airflow/airflow/cli/cli_parser.py
index afa92246d..c77d279e5 100644
--- a/flink-ai-flow/lib/airflow/airflow/cli/cli_parser.py
+++ b/flink-ai-flow/lib/airflow/airflow/cli/cli_parser.py
@@ -1437,6 +1437,7 @@ class GroupCommand(NamedTuple):
             ARG_STDOUT,
             ARG_STDERR,
             ARG_LOG_FILE,
+            ARG_SERVER_URI,
         ),
         epilog=(
             'Signals:\n'
diff --git a/flink-ai-flow/lib/airflow/airflow/cli/commands/scheduler_command.py b/flink-ai-flow/lib/airflow/airflow/cli/commands/scheduler_command.py
index 100a0f1d0..e15058754 100644
--- a/flink-ai-flow/lib/airflow/airflow/cli/commands/scheduler_command.py
+++ b/flink-ai-flow/lib/airflow/airflow/cli/commands/scheduler_command.py
@@ -19,6 +19,8 @@
 import signal
 
 import daemon
+from airflow.contrib.jobs.event_based_scheduler_job import EventBasedSchedulerJob
+from airflow.contrib.jobs.scheduler_factory import SchedulerFactory
 from daemon.pidfile import TimeoutPIDLockFile
 
 from airflow import settings
@@ -36,6 +38,16 @@ def scheduler(args):
         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(
diff --git a/flink-ai-flow/lib/airflow/airflow/config_templates/default_airflow.cfg b/flink-ai-flow/lib/airflow/airflow/config_templates/default_airflow.cfg
index 8250fe5a4..b3aa8af08 100644
--- a/flink-ai-flow/lib/airflow/airflow/config_templates/default_airflow.cfg
+++ b/flink-ai-flow/lib/airflow/airflow/config_templates/default_airflow.cfg
@@ -47,6 +47,10 @@ hostname_callable = socket.getfqdn
 # can be utc (default), system, or any IANA timezone string (e.g. Europe/Amsterdam)
 default_timezone = utc
 
+# The Scheduler Job class that airflow should use. Choie include
+# ``SchedulerJob``, ``EventBasedSchedulerJob``
+# scheduler =
+
 # The executor class that airflow should use. Choices include
 # ``SequentialExecutor``, ``LocalExecutor``, ``CeleryExecutor``, ``DaskExecutor``,
 # ``KubernetesExecutor``, ``CeleryKubernetesExecutor`` or the
diff --git a/flink-ai-flow/lib/airflow/airflow/contrib/jobs/event_based_scheduler_job.py b/flink-ai-flow/lib/airflow/airflow/contrib/jobs/event_based_scheduler_job.py
index ee67baeea..db740553a 100644
--- a/flink-ai-flow/lib/airflow/airflow/contrib/jobs/event_based_scheduler_job.py
+++ b/flink-ai-flow/lib/airflow/airflow/contrib/jobs/event_based_scheduler_job.py
@@ -110,74 +110,73 @@ def repeat(*args, **kwargs):
         )
         self.timers.run()
 
-    def _stop_timer(self):
+    def stop_timer(self):
         if self.timers and self._timer_handler:
             self.timers.cancel(self._timer_handler)
 
     def submit_sync_thread(self):
         threading.Thread(target=self.sync).start()
 
-    def schedule(self):
-        self.log.info("Starting the scheduler.")
-        self._restore_unfinished_dag_run()
-        while True:
-            identified_message = self.mailbox.get_identified_message()
-            origin_event = identified_message.deserialize()
-            self.log.debug("Event: {}".format(origin_event))
-            if SchedulerInnerEventUtil.is_inner_event(origin_event):
-                event = SchedulerInnerEventUtil.to_inner_event(origin_event)
+    def schedule(self) -> bool:
+        identified_message = self.mailbox.get_identified_message()
+        if not identified_message:
+            return True
+        origin_event = identified_message.deserialize()
+        self.log.debug("Event: {}".format(origin_event))
+        if SchedulerInnerEventUtil.is_inner_event(origin_event):
+            event = SchedulerInnerEventUtil.to_inner_event(origin_event)
+        else:
+            event = origin_event
+        with create_session() as session:
+            if isinstance(event, BaseEvent):
+                dagruns = self._find_dagruns_by_event(event, session)
+                for dagrun in dagruns:
+                    dag_run_id = DagRunId(dagrun.dag_id, dagrun.run_id)
+                    self.task_event_manager.handle_event(dag_run_id, event)
+            elif isinstance(event, RequestEvent):
+                self._process_request_event(event)
+            elif isinstance(event, TaskSchedulingEvent):
+                self._schedule_task(event)
+            elif isinstance(event, TaskStatusChangedEvent):
+                dagrun = self._find_dagrun(event.dag_id, event.execution_date, session)
+                tasks = self._find_scheduled_tasks(dagrun, session)
+                self._send_scheduling_task_events(tasks, SchedulingAction.START)
+                if dagrun.state in State.finished:
+                    self.mailbox.send_message(DagRunFinishedEvent(dagrun.run_id).to_event())
+            elif isinstance(event, DagExecutableEvent):
+                dagrun = self._create_dag_run(event.dag_id, session=session)
+                tasks = self._find_scheduled_tasks(dagrun, session)
+                self._send_scheduling_task_events(tasks, SchedulingAction.START)
+            elif isinstance(event, EventHandleEvent):
+                dag_runs = DagRun.find(dag_id=event.dag_id, run_id=event.dag_run_id)
+                assert len(dag_runs) == 1
+                ti = dag_runs[0].get_task_instance(event.task_id)
+                self._send_scheduling_task_event(ti, event.action)
+            elif isinstance(event, StopDagEvent):
+                self._stop_dag(event.dag_id, session)
+            elif isinstance(event, DagRunFinishedEvent):
+                self._remove_periodic_events(event.run_id)
+            elif isinstance(event, PeriodicEvent):
+                dag_runs = DagRun.find(run_id=event.run_id)
+                assert len(dag_runs) == 1
+                ti = dag_runs[0].get_task_instance(event.task_id)
+                self._send_scheduling_task_event(ti, SchedulingAction.RESTART)
+            elif isinstance(event, StopSchedulerEvent):
+                self.log.info("{} {}".format(self.id, event.job_id))
+                if self.id == event.job_id or 0 == event.job_id:
+                    self.log.info("break the scheduler event loop.")
+                    identified_message.remove_handled_message()
+                    session.expunge_all()
+                    return False
+            elif isinstance(event, ParseDagRequestEvent) or isinstance(event, ParseDagResponseEvent):
+                pass
+            elif isinstance(event, ResponseEvent):
+                pass
             else:
-                event = origin_event
-            with create_session() as session:
-                if isinstance(event, BaseEvent):
-                    dagruns = self._find_dagruns_by_event(event, session)
-                    for dagrun in dagruns:
-                        dag_run_id = DagRunId(dagrun.dag_id, dagrun.run_id)
-                        self.task_event_manager.handle_event(dag_run_id, event)
-                elif isinstance(event, RequestEvent):
-                    self._process_request_event(event)
-                elif isinstance(event, TaskSchedulingEvent):
-                    self._schedule_task(event)
-                elif isinstance(event, TaskStatusChangedEvent):
-                    dagrun = self._find_dagrun(event.dag_id, event.execution_date, session)
-                    tasks = self._find_scheduled_tasks(dagrun, session)
-                    self._send_scheduling_task_events(tasks, SchedulingAction.START)
-                    if dagrun.state in State.finished:
-                        self.mailbox.send_message(DagRunFinishedEvent(dagrun.run_id).to_event())
-                elif isinstance(event, DagExecutableEvent):
-                    dagrun = self._create_dag_run(event.dag_id, session=session)
-                    tasks = self._find_scheduled_tasks(dagrun, session)
-                    self._send_scheduling_task_events(tasks, SchedulingAction.START)
-                elif isinstance(event, EventHandleEvent):
-                    dag_runs = DagRun.find(dag_id=event.dag_id, run_id=event.dag_run_id)
-                    assert len(dag_runs) == 1
-                    ti = dag_runs[0].get_task_instance(event.task_id)
-                    self._send_scheduling_task_event(ti, event.action)
-                elif isinstance(event, StopDagEvent):
-                    self._stop_dag(event.dag_id, session)
-                elif isinstance(event, DagRunFinishedEvent):
-                    self._remove_periodic_events(event.run_id)
-                elif isinstance(event, PeriodicEvent):
-                    dag_runs = DagRun.find(run_id=event.run_id)
-                    assert len(dag_runs) == 1
-                    ti = dag_runs[0].get_task_instance(event.task_id)
-                    self._send_scheduling_task_event(ti, SchedulingAction.RESTART)
-                elif isinstance(event, StopSchedulerEvent):
-                    self.log.info("{} {}".format(self.id, event.job_id))
-                    if self.id == event.job_id or 0 == event.job_id:
-                        self.log.info("break the scheduler event loop.")
-                        identified_message.remove_handled_message()
-                        session.expunge_all()
-                        break
-                elif isinstance(event, ParseDagRequestEvent) or isinstance(event, ParseDagResponseEvent):
-                    pass
-                elif isinstance(event, ResponseEvent):
-                    pass
-                else:
-                    self.log.error("can not handler the event {}".format(event))
-                identified_message.remove_handled_message()
-                session.expunge_all()
-        self._stop_timer()
+                self.log.error("can not handler the event {}".format(event))
+            identified_message.remove_handled_message()
+            session.expunge_all()
+            return True
 
     def stop(self) -> None:
         self.mailbox.send_message(StopSchedulerEvent(self.id).to_event())
@@ -479,7 +478,7 @@ def _reset_unfinished_task_state(dag_run):
             session.commit()
 
     @provide_session
-    def _restore_unfinished_dag_run(self, session):
+    def restore_unfinished_dag_run(self, session):
         dag_runs = DagRun.next_dagruns_to_examine(session, max_number=sys.maxsize).all()
         if not dag_runs or len(dag_runs) == 0:
             return
@@ -626,7 +625,7 @@ def _execute(self):
 
             self.scheduler.submit_sync_thread()
             self.scheduler.recover(self.last_scheduling_id)
-            self.scheduler.schedule()
+            self._run_scheduler_loop()
 
             self.executor.end()
             self.periodic_manager.shutdown()
@@ -640,6 +639,15 @@ def _execute(self):
         finally:
             self.log.info("Exited execute loop")
 
+    def _run_scheduler_loop(self) -> None:
+        self.log.info("Starting the scheduler loop.")
+        self.scheduler.restore_unfinished_dag_run()
+        should_continue = True
+        while should_continue:
+            should_continue = self.scheduler.schedule()
+            self.heartbeat(only_if_necessary=True)
+        self.scheduler.stop_timer()
+
     def _start_listen_events(self):
         watcher = SchedulerEventWatcher(self.mailbox)
         self.notification_client.start_listen_events(
diff --git a/flink-ai-flow/lib/airflow/airflow/contrib/jobs/scheduler_factory.py b/flink-ai-flow/lib/airflow/airflow/contrib/jobs/scheduler_factory.py
new file mode 100644
index 000000000..dad441bd4
--- /dev/null
+++ b/flink-ai-flow/lib/airflow/airflow/contrib/jobs/scheduler_factory.py
@@ -0,0 +1,46 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import logging
+
+from airflow.configuration import conf
+from airflow.utils.module_loading import import_string
+
+log = logging.getLogger(__name__)
+
+
+class SchedulerFactory:
+
+    DEFAULT_SCHEDULER = "SchedulerJob"
+    EVENT_BASED_SCHEDULER = "EventBasedSchedulerJob"
+
+    executors = {
+        DEFAULT_SCHEDULER: 'airflow.jobs.scheduler_job.SchedulerJob',
+        EVENT_BASED_SCHEDULER: 'airflow.contrib.jobs.event_based_scheduler_job.EventBasedSchedulerJob',
+    }
+
+    @classmethod
+    def get_default_scheduler(cls):
+        scheduler_name = cls.get_scheduler_name()
+        log.debug("Loading core scheduler: %s", scheduler_name)
+        if scheduler_name in cls.executors:
+            return import_string(cls.executors[scheduler_name])
+        else:
+            return import_string(scheduler_name)
+
+    @classmethod
+    def get_scheduler_name(cls):
+        return conf.get('core', 'scheduler', fallback=cls.DEFAULT_SCHEDULER)
diff --git a/flink-ai-flow/lib/airflow/airflow/utils/mailbox.py b/flink-ai-flow/lib/airflow/airflow/utils/mailbox.py
index 80a3f060b..34d8020ce 100644
--- a/flink-ai-flow/lib/airflow/airflow/utils/mailbox.py
+++ b/flink-ai-flow/lib/airflow/airflow/utils/mailbox.py
@@ -17,6 +17,7 @@
 import pickle
 import queue
 
+from airflow.configuration import conf
 from airflow.models.message import Message, IdentifiedMessage
 from airflow.utils.log.logging_mixin import LoggingMixin
 
@@ -46,7 +47,7 @@ def get_message(self):
             return None
 
     def get_identified_message(self) -> IdentifiedMessage:
-        return self.queue.get()
+        return self.get_message_with_timeout(timeout=1)
 
     def length(self):
         return self.queue.qsize()
diff --git a/flink-ai-flow/lib/airflow/airflow/www/views.py b/flink-ai-flow/lib/airflow/airflow/www/views.py
index 150c26ad9..75262425f 100644
--- a/flink-ai-flow/lib/airflow/airflow/www/views.py
+++ b/flink-ai-flow/lib/airflow/airflow/www/views.py
@@ -35,6 +35,7 @@
 import nvd3
 import sqlalchemy as sqla
 import yaml
+from airflow.contrib.jobs.scheduler_factory import SchedulerFactory
 from flask import (
     Markup,
     Response,
@@ -427,10 +428,11 @@ class AirflowBaseView(BaseView):  # noqa: D101
     }
 
     def render_template(self, *args, **kwargs):
+        scheduler_class = SchedulerFactory.get_default_scheduler()
         return super().render_template(
             *args,
             # Cache this at most once per request, not for the lifetime of the view instance
-            scheduler_job=lazy_object_proxy.Proxy(SchedulerJob.most_recent_job),
+            scheduler_job=lazy_object_proxy.Proxy(scheduler_class.most_recent_job),
             **kwargs,
         )
 
@@ -456,7 +458,8 @@ def health(self):
         scheduler_status = 'unhealthy'
         payload['metadatabase'] = {'status': 'healthy'}
         try:
-            scheduler_job = SchedulerJob.most_recent_job()
+            scheduler_class = SchedulerFactory.get_default_scheduler()
+            scheduler_job = scheduler_class.most_recent_job()
 
             if scheduler_job:
                 latest_scheduler_heartbeat = scheduler_job.latest_heartbeat.isoformat()


Test patch:
diff --git a/flink-ai-flow/lib/airflow/tests/contrib/jobs/test_scheduler_factory.py b/flink-ai-flow/lib/airflow/tests/contrib/jobs/test_scheduler_factory.py
new file mode 100644
index 000000000..44ee369f4
--- /dev/null
+++ b/flink-ai-flow/lib/airflow/tests/contrib/jobs/test_scheduler_factory.py
@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import unittest
+
+from airflow.configuration import conf
+from airflow.contrib.jobs.event_based_scheduler_job import EventBasedSchedulerJob
+from airflow.contrib.jobs.scheduler_factory import SchedulerFactory
+from airflow.jobs.scheduler_job import SchedulerJob
+
+
+class TestSchedulerFactory(unittest.TestCase):
+    def test_get_default_scheduler(self):
+        conf.set('core', 'scheduler', SchedulerFactory.DEFAULT_SCHEDULER)
+        scheduler_class = SchedulerFactory.get_default_scheduler()
+        self.assertEqual(scheduler_class, SchedulerJob)
+
+        conf.set('core', 'scheduler', SchedulerFactory.EVENT_BASED_SCHEDULER)
+        scheduler_class = SchedulerFactory.get_default_scheduler()
+        self.assertEqual(scheduler_class, EventBasedSchedulerJob)