uranusjr commented on code in PR #74789:
URL: https://github.com/apache/airflow/pull/74789#discussion_r4344022728
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -1509,17 +1598,61 @@ def _create_process(self, dag_file: DagFileInfo) ->
BaseDagFileProcessorProcess:
client=self.client,
)
+ def _get_dispatch_source(self, file: DagFileInfo) -> DagParseSource | None:
+ if file in self._callback_to_execute and (
+ file.bundle_version is not None or file.bundle_name not in
self._bundles_waiting_for_refresh
+ ):
+ # Callbacks can run after refresh failures, but pending refreshes
must drain unversioned callbacks.
+ return DagParseSource(bundle_version=file.bundle_version)
+ if (
+ file.bundle_name in self._bundles_waiting_for_refresh
+ or file.bundle_name in self._force_refresh_bundles
+ ):
+ return None
+ return self._bundle_parse_sources.get(file.bundle_name)
+
+ def _readmit_files_with_source(self) -> None:
+ ready = [file for file in self._files_awaiting_source if
self._get_dispatch_source(file) is not None]
+ if not ready:
+ return
+ for file in ready:
+ del self._files_awaiting_source[file]
+
+ queued_first = dict.fromkeys([*self._file_queue, *ready])
+ restored_first = dict.fromkeys([*ready, *self._file_queue])
+ callbacks = [file for file in queued_first if file in
self._callback_to_execute]
+ priority = [
+ file
+ for file in queued_first
+ if file in self._priority_files and file not in
self._callback_to_execute
+ ]
+ regular = [
+ file
+ for file in restored_first
+ if file not in self._priority_files and file not in
self._callback_to_execute
+ ]
+ self._file_queue.clear()
+ self._add_files_to_queue(callbacks + priority + regular, mode="back")
+
def _start_new_processes(self):
"""Start more processors if we have enough slots and files to
process."""
+ self._readmit_files_with_source()
bundle_to_team = self._get_team_names({file.bundle_name for file in
self._file_queue})
while self._parallelism > len(self._processors) and self._file_queue:
file, _ = self._file_queue.popitem(last=False)
# Stop creating duplicate processor i.e. processor with the same
filepath
if file in self._processors:
+ self._priority_files.discard(file)
continue
+ source = self._get_dispatch_source(file)
+ if source is None:
+ self._files_awaiting_source[file] = None
+ continue
processor = self._create_process(file)
+ processor.parse_source = source
Review Comment:
I feel it’d be cleaner to set this in `_create_process`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
