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

Issue:
[Flink AI Flow][Bug] Failed to trigger dagrun cause of dag parse not finish



Discussion / hints:

Fix patch:
diff --git a/flink-ai-flow/ai_flow/api/execution.py b/flink-ai-flow/ai_flow/api/execution.py
index 5ea83f06a..52ff9def4 100644
--- a/flink-ai-flow/ai_flow/api/execution.py
+++ b/flink-ai-flow/ai_flow/api/execution.py
@@ -85,14 +85,19 @@ def resume_workflow(workflow_name) -> bool:
         """
         pass
 
-    def trigger_workflow_execution(self, workflow_name) -> ExecutionContext:
+    def trigger_workflow_execution(self, project_desc, workflow_name) -> ExecutionContext:
         """
         Trigger a new instance of workflow immediately.
 
         :param workflow_name: workflow name
+        :param project_desc: project desc
         :return: True if a new instance is triggered
         """
-        self.airflow_client.trigger_parse_dag()
+        deploy_path = project_desc.project_config.get_airflow_deploy_path()
+        if deploy_path is None:
+            raise Exception("airflow_deploy_path config not set!")
+        airflow_file_path = deploy_path + '/' + workflow_name + '.py'
+        self.airflow_client.trigger_parse_dag(airflow_file_path)
         return self.airflow_client.schedule_dag(workflow_name)
 
     def stop_workflow_execution(self, workflow_name, context) -> bool:
@@ -160,6 +165,6 @@ def restart_task_instance(self, workflow_name, job_name, context: ExecutionConte
         else:
             return False
 
-    def trigger_parse_dag(self):
-        """Trigger a dag parse. """
-        return self.airflow_client.trigger_parse_dag()
+    def trigger_parse_dag(self, file_path):
+        """Trigger a dag parse of specific file. """
+        return self.airflow_client.trigger_parse_dag(file_path)
diff --git a/flink-ai-flow/ai_flow/api/project.py b/flink-ai-flow/ai_flow/api/project.py
index 480e76747..8a2364ec0 100644
--- a/flink-ai-flow/ai_flow/api/project.py
+++ b/flink-ai-flow/ai_flow/api/project.py
@@ -224,7 +224,7 @@ def deploy_to_airflow(project_path: Text = None,
 def _submit_to_airflow(project_desc: ProjectDesc = ProjectDesc(), dag_id: Text = None):
     airflow_operation = AirflowOperation(
         notification_server_uri=project_desc.project_config.get_notification_service_uri())
-    return airflow_operation.trigger_workflow_execution(dag_id)
+    return airflow_operation.trigger_workflow_execution(project_desc, dag_id)
 
 
 def generate_project_desc(project_path: Text = None) -> ProjectDesc:
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 05da97a81..ee67baeea 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
@@ -564,7 +564,7 @@ class EventBasedSchedulerJob(BaseJob):
     def __init__(self, dag_directory,
                  server_uri=None,
                  max_runs=-1,
-                 refresh_dag_dir_interval=conf.getint('scheduler', 'refresh_dag_dir_interval', fallback=30),
+                 refresh_dag_dir_interval=conf.getint('scheduler', 'refresh_dag_dir_interval', fallback=1),
                  *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.mailbox: Mailbox = Mailbox()
diff --git a/flink-ai-flow/lib/airflow/airflow/contrib/jobs/scheduler_client.py b/flink-ai-flow/lib/airflow/airflow/contrib/jobs/scheduler_client.py
index 9a5f2d5aa..08157f105 100644
--- a/flink-ai-flow/lib/airflow/airflow/contrib/jobs/scheduler_client.py
+++ b/flink-ai-flow/lib/airflow/airflow/contrib/jobs/scheduler_client.py
@@ -52,7 +52,7 @@ def __init__(self, server_uri=None, namespace=None, ns_client=None):
     def generate_id(id):
         return '{}_{}'.format(id, time.time_ns())
 
-    def trigger_parse_dag(self) -> bool:
+    def trigger_parse_dag(self, file_path) -> bool:
         id = self.generate_id('')
         watcher: ResponseWatcher = ResponseWatcher()
         handler: ThreadEventWatcherHandle \
@@ -61,7 +61,8 @@ def trigger_parse_dag(self) -> bool:
                                                 namespace=SCHEDULER_NAMESPACE, watcher=watcher)
 
         self.ns_client.send_event(BaseEvent(key=id,
-                                            event_type=SchedulerInnerEventType.PARSE_DAG_REQUEST.value, value=''))
+                                            event_type=SchedulerInnerEventType.PARSE_DAG_REQUEST.value,
+                                            value=file_path))
         result = watcher.get_result()
         handler.stop()
         return True
diff --git a/flink-ai-flow/lib/airflow/airflow/utils/dag_processing.py b/flink-ai-flow/lib/airflow/airflow/utils/dag_processing.py
index ab577ebfa..af47c3a76 100644
--- a/flink-ai-flow/lib/airflow/airflow/utils/dag_processing.py
+++ b/flink-ai-flow/lib/airflow/airflow/utils/dag_processing.py
@@ -53,7 +53,8 @@
 from airflow.utils.process_utils import kill_child_processes_by_pids, reap_process_group
 from airflow.utils.session import provide_session
 from airflow.utils.state import State
-from airflow.events.scheduler_events import SCHEDULER_NAMESPACE
+from airflow.events.scheduler_events import SCHEDULER_NAMESPACE, SchedulerInnerEventType
+
 
 class AbstractDagFileProcessorProcess(metaclass=ABCMeta):
     """Processes a DAG file. See SchedulerJob.process_file() for more details."""
@@ -478,7 +479,7 @@ def __init__(self, signal_queue: queue.Queue):
     def process(self, events: List[BaseEvent]):
         for e in events:
             self.log.debug("Listen Event: {}".format(e))
-            self.signal_queue.put(e)
+            self.signal_queue.put((e, timezone.utcnow()))
 
 
 class DagFileProcessorManager(LoggingMixin):  # pylint: disable=too-many-instance-attributes
@@ -540,7 +541,7 @@ def __init__(
         self.signal_queue = queue.Queue()
         if notification_service_uri is not None:
             self.watcher = ProcessorManagerWatcher(self.signal_queue)
-            self.message_buffer = []
+            self.message_buffer: Dict[str, (BaseEvent, datetime)] = {}
 
         self._parallelism = conf.getint('scheduler', 'parsing_processes')
         if 'sqlite' in conf.get('core', 'sql_alchemy_conn') and self._parallelism > 1:
@@ -759,31 +760,41 @@ def _run_parsing_loop(self):
                 else:
                     poll_time = 0.0
             # send event to notify parse dag finished
+            refresh_dag_dir_interval = time.monotonic() - loop_start_time
+            if refresh_dag_dir_interval < self._refresh_dag_dir_interval:
+                wait_time = self._refresh_dag_dir_interval - refresh_dag_dir_interval
+            else:
+                wait_time = 1.0
             if self.notification_service_uri is not None and len(self.message_buffer) > 0:
-                for e in self.message_buffer:
-                    self.log.debug('send message key {}'.format(e.key))
-                    self.ns_client.send_event(BaseEvent(key=e.key, value='', event_type='PARSE_DAG_RESPONSE'))
-                self.message_buffer.clear()
-
-            if 0 != self._refresh_dag_dir_interval:
-                refresh_dag_dir_interval = time.monotonic() - loop_start_time
-                if refresh_dag_dir_interval < self._refresh_dag_dir_interval:
-                    wait_time = self._refresh_dag_dir_interval - refresh_dag_dir_interval
-                else:
-                    wait_time = 1.0
-                self.log.info('Dag ProcessorManager sleep {}.'
-                              .format(wait_time))
-                try:
-                    message: BaseEvent = self.signal_queue.get(timeout=wait_time)
-                    self.log.debug('receive message key {}'.format(message.key))
-                    self.message_buffer.append(message)
-                    if self.signal_queue.qsize() > 0:
-                        for i in range(self.signal_queue.qsize()):
-                            message = self.signal_queue.get()
-                            self.log.debug('receive message key {}'.format(message.key))
-                            self.message_buffer.append(message)
-                except Exception as e:
-                    pass
+                self._process_and_send_response()
+                self.collect_results()
+                time.sleep(wait_time)
+
+            if self.signal_queue.qsize() > 0:
+                self._add_file_to_queue()
+
+    def _process_and_send_response(self):
+        for event, process_time in self.message_buffer.copy().values():
+            file_path = event.value
+            finish_process_time = self.get_last_finish_time(file_path)
+            duration = self.get_last_runtime(file_path)
+            if not finish_process_time or not duration:
+                continue
+            start_process_time = finish_process_time - timedelta(seconds=duration)
+            if start_process_time > process_time:
+                self.ns_client.send_event(BaseEvent(key=event.key,
+                                                    value=file_path,
+                                                    event_type=SchedulerInnerEventType.PARSE_DAG_RESPONSE.value))
+                self.message_buffer.pop(file_path)
+
+    def _add_file_to_queue(self):
+        for i in range(self.signal_queue.qsize()):
+            message, process_time = self.signal_queue.get()
+            file_path = message.value
+            if file_path in self._file_path_queue:
+                self._file_path_queue.remove(file_path)
+            self._file_path_queue.insert(0, file_path)
+            self.message_buffer[file_path] = (message, process_time)
 
     def _add_callback_to_queue(self, request: CallbackRequest):
         self._callback_to_execute[request.full_filepath].append(request)


Test patch:
diff --git a/flink-ai-flow/lib/airflow/tests/contrib/jobs/test_dag_trigger.py b/flink-ai-flow/lib/airflow/tests/contrib/jobs/test_dag_trigger.py
index 35b60506f..adc6ecdb2 100644
--- a/flink-ai-flow/lib/airflow/tests/contrib/jobs/test_dag_trigger.py
+++ b/flink-ai-flow/lib/airflow/tests/contrib/jobs/test_dag_trigger.py
@@ -14,22 +14,26 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import multiprocessing
 import time
 import unittest
 
 import pendulum.datetime
+from notification_service.base_notification import BaseEvent
+from notification_service.client import NotificationClient
 from notification_service.event_storage import MemoryEventStorage
 from notification_service.master import NotificationMaster
 from notification_service.service import NotificationService
 from airflow.models.serialized_dag import SerializedDagModel
 
 from airflow.contrib.jobs.dag_trigger import DagTrigger
-from airflow.contrib.jobs.scheduler_client import EventSchedulerClient
+from airflow.contrib.jobs.scheduler_client import EventSchedulerClient, ResponseWatcher
 from airflow.models import DagModel
 from airflow.utils.mailbox import Mailbox
 from airflow.utils.session import create_session
-from airflow.events.scheduler_events import SchedulerInnerEventUtil
+from airflow.events.scheduler_events import SchedulerInnerEventUtil, SchedulerInnerEventType, SCHEDULER_NAMESPACE
 from tests.test_utils import db
+from tests.test_utils.config import conf_vars
 
 
 class TestDagTrigger(unittest.TestCase):
@@ -106,6 +110,44 @@ def test_file_processor_manager_kill(self):
         assert dag_file_processor_manager_process.is_alive()
         dag_trigger.end()
 
+    @conf_vars({('core', 'load_examples'): 'False'})
+    def test_trigger_parse_dag(self):
+        import os
+        port = 50102
+        server_uri = "localhost:{}".format(port)
+        storage = MemoryEventStorage()
+        master = NotificationMaster(NotificationService(storage), port)
+        master.run()
+        dag_folder = os.path.abspath(os.path.dirname(__file__)) + "/../../dags"
+        mailbox = Mailbox()
+        dag_trigger = DagTrigger(dag_folder, -1, [], False, mailbox, notification_service_uri=server_uri)
+        dag_trigger.start()
+
+        to_be_triggered = [dag_folder + "/test_event_based_scheduler.py",
+                           dag_folder + "/test_event_task_dag.py",
+                           dag_folder + "/test_event_based_executor.py",
+                           dag_folder + "/test_scheduler_dags.py",
+                           ]
+        for file in to_be_triggered:
+            self._send_request_and_receive_response(server_uri, file)
+        dag_trigger.end()
+
+    def _send_request_and_receive_response(self, server_uri, file_path):
+        key = '{}_{}'.format(file_path, time.time_ns())
+        client = NotificationClient(server_uri=server_uri,
+                                    default_namespace=SCHEDULER_NAMESPACE)
+        event = BaseEvent(key=key,
+                          event_type=SchedulerInnerEventType.PARSE_DAG_REQUEST.value,
+                          value=file_path)
+        client.send_event(event)
+        watcher: ResponseWatcher = ResponseWatcher()
+        client.start_listen_event(key=key,
+                                  event_type=SchedulerInnerEventType.PARSE_DAG_RESPONSE.value,
+                                  watcher=watcher)
+        res: BaseEvent = watcher.get_result()
+        self.assertEquals(event.key, res.key)
+        self.assertEquals(event.value, file_path)
+
     @staticmethod
     def _add_dag_needing_dagrun():
         with create_session() as session:
diff --git a/flink-ai-flow/lib/airflow/tests/dags/test_event_based_executor.py b/flink-ai-flow/lib/airflow/tests/dags/test_event_based_executor.py
index 55eadaac9..5a62fe280 100644
--- a/flink-ai-flow/lib/airflow/tests/dags/test_event_based_executor.py
+++ b/flink-ai-flow/lib/airflow/tests/dags/test_event_based_executor.py
@@ -24,7 +24,7 @@
 dag = DAG(
     dag_id='test_event_based_dag',
     default_args=default_args,
-    schedule_interval=timedelta(days=10000),
+    schedule_interval=timedelta(days=1),
 )
 
 t1 = BashOperator(