jedcunningham commented on code in PR #30255:
URL: https://github.com/apache/airflow/pull/30255#discussion_r1161135548
##########
airflow/jobs/base_job.py:
##########
@@ -276,5 +251,44 @@ def run(self) -> int | None:
Stats.incr(self.__class__.__name__.lower() + "_end", 1, 1)
return ret
- def _execute(self) -> int | None:
- raise NotImplementedError("This method needs to be overridden")
+ @property
+ def job_runner(self) -> BaseJobRunner:
+ """Returns the job runner instance."""
+ return self._job_runner
+
+ @provide_session
+ def most_recent_job(self, session=NEW_SESSION) -> BaseJob | None:
+ """Returns the most recent job of this type, if any, based on last
heartbeat received."""
+ return most_recent_job(self.job_type, session=session)
+
+
+@provide_session
+def perform_heartbeat(job: BaseJob | BaseJobPydantic, only_if_necessary: bool,
session=None):
Review Comment:
```suggestion
def perform_heartbeat(job: BaseJob | BaseJobPydantic, only_if_necessary:
bool, session: Session = NEW_SESSION):
```
##########
airflow/jobs/scheduler_job.py:
##########
@@ -161,12 +165,16 @@ def __init__(
# Check what SQL backend we use
sql_conn: str = conf.get_mandatory_value("database",
"sql_alchemy_conn").lower()
self.using_sqlite = sql_conn.startswith("sqlite")
- # Dag Processor agent - not used in Dag Processor standalone mode.
+ # Dag Processor agent - not used in Dag Processor eartstandalone mode.
Review Comment:
```suggestion
# Dag Processor agent - not used in Dag Processor standalone mode.
```
I figure this wasn't intentional.
##########
airflow/www/views.py:
##########
@@ -635,12 +635,14 @@ class AirflowBaseView(BaseView):
def render_template(self, *args, **kwargs):
# Add triggerer_job only if we need it
- if TriggererJob.is_needed():
- kwargs["triggerer_job"] =
lazy_object_proxy.Proxy(TriggererJob.most_recent_job)
+ if TriggererJobRunner.is_needed():
+ kwargs["triggerer_job"] = lazy_object_proxy.Proxy(
+ lambda: most_recent_job(TriggererJobRunner.job_type)
+ )
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(lambda:
most_recent_job(SchedulerJobRunner.job_type)),
Review Comment:
```suggestion
scheduler_job=lazy_object_proxy.Proxy(SchedulerJobRunner.most_recent_job),
```
This one too?
##########
airflow/dag_processing/manager.py:
##########
@@ -576,8 +585,11 @@ def _run_parsing_loop(self):
while True:
loop_start_time = time.monotonic()
ready = multiprocessing.connection.wait(self.waitables.keys(),
timeout=poll_time)
- if self._job:
- self._job.heartbeat()
+ # we cannot (for now) define job in _job_runner nicely due to
circular references of
+ # job and job runner, so we have to use hasattr, but we might
address it in the future
Review Comment:
nit: comment says hasattr, but you are using getattr
##########
airflow/jobs/base_job.py:
##########
Review Comment:
Given this new architecture, I'm not sure I like the BaseJob name any
longer. Did you consider changing it?
##########
airflow/jobs/job_runner.py:
##########
@@ -0,0 +1,48 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.utils.session import NEW_SESSION, provide_session
+
+if TYPE_CHECKING:
+ from airflow.jobs.base_job import BaseJob
Review Comment:
```suggestion
from sqlalchemy.orm import Session
from airflow.jobs.base_job import BaseJob
```
##########
airflow/jobs/base_job.py:
##########
@@ -276,5 +251,44 @@ def run(self) -> int | None:
Stats.incr(self.__class__.__name__.lower() + "_end", 1, 1)
return ret
- def _execute(self) -> int | None:
- raise NotImplementedError("This method needs to be overridden")
+ @property
+ def job_runner(self) -> BaseJobRunner:
+ """Returns the job runner instance."""
+ return self._job_runner
+
+ @provide_session
+ def most_recent_job(self, session=NEW_SESSION) -> BaseJob | None:
Review Comment:
```suggestion
def most_recent_job(self, session: Session = NEW_SESSION) -> BaseJob |
None:
```
##########
airflow/jobs/job_runner.py:
##########
@@ -0,0 +1,48 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.utils.session import NEW_SESSION, provide_session
+
+if TYPE_CHECKING:
+ from airflow.jobs.base_job import BaseJob
+ from airflow.serialization.pydantic.base_job import BaseJobPydantic
+
+
+class BaseJobRunner:
+ """Abstract class for job runners to derive from."""
+
+ job_type = "undefined"
+
+ job: BaseJob | BaseJobPydantic
+
+ def _execute(self) -> int | None:
+ """
+ Executes the logic connected to the runner. This method should be
+ overridden by subclasses.
+
+ :meta private:
+ :return: return code if available, otherwise None
+ """
+ raise NotImplementedError()
+
+ @provide_session
+ def heartbeat_callback(self, session=NEW_SESSION) -> None:
Review Comment:
```suggestion
def heartbeat_callback(self, session: Session = NEW_SESSION) -> None:
```
##########
airflow/jobs/scheduler_job.py:
##########
@@ -161,12 +165,16 @@ def __init__(
# Check what SQL backend we use
sql_conn: str = conf.get_mandatory_value("database",
"sql_alchemy_conn").lower()
self.using_sqlite = sql_conn.startswith("sqlite")
- # Dag Processor agent - not used in Dag Processor standalone mode.
+ # Dag Processor agent - not used in Dag Processor eartstandalone mode.
self.processor_agent: DagFileProcessorAgent | None = None
self.dagbag = DagBag(dag_folder=self.subdir, read_dags_from_db=True,
load_op_links=False)
self._paused_dag_without_running_dagruns: set = set()
+ @provide_session
+ def heartbeat_callback(self, session=NEW_SESSION) -> None:
Review Comment:
```suggestion
def heartbeat_callback(self, session: Session = NEW_SESSION) -> None:
```
##########
airflow/jobs/local_task_job.py:
##########
@@ -293,5 +295,6 @@ def heartbeat_callback(self, session: Session =
NEW_SESSION) -> None:
def _log_return_code_metric(self, return_code: int):
Stats.incr(
-
f"local_task_job.task_exit.{self.id}.{self.dag_id}.{self.task_instance.task_id}.{return_code}"
+ f"local_task_job.task_exit."
Review Comment:
```suggestion
"local_task_job.task_exit."
```
##########
airflow/models/skipmixin.py:
##########
@@ -143,7 +146,12 @@ def skip(
session=session,
)
- def skip_all_except(self, ti: TaskInstance, branch_task_ids: None | str |
Iterable[str]):
+ def skip_all_except(
+ self,
+ ti: TaskInstance | TaskInstancePydantic,
+ task: BaseOperator,
Review Comment:
Hmm, this feels like it could be a breaking change? Why do we need to do
this?
##########
airflow/jobs/base_job.py:
##########
@@ -276,5 +251,44 @@ def run(self) -> int | None:
Stats.incr(self.__class__.__name__.lower() + "_end", 1, 1)
return ret
- def _execute(self) -> int | None:
- raise NotImplementedError("This method needs to be overridden")
+ @property
+ def job_runner(self) -> BaseJobRunner:
+ """Returns the job runner instance."""
+ return self._job_runner
+
+ @provide_session
+ def most_recent_job(self, session=NEW_SESSION) -> BaseJob | None:
+ """Returns the most recent job of this type, if any, based on last
heartbeat received."""
+ return most_recent_job(self.job_type, session=session)
+
+
+@provide_session
+def perform_heartbeat(job: BaseJob | BaseJobPydantic, only_if_necessary: bool,
session=None):
+ if isinstance(job, BaseJob):
+ job.heartbeat()
+ else:
+ # TODO (potiuk): Make it works over internal API as a follow up
+ BaseJob.get(job.id,
session=session).heartbeat(only_if_necessary=only_if_necessary)
+
+
+@provide_session
+def most_recent_job(job_type: str, session=NEW_SESSION) -> BaseJob | None:
Review Comment:
```suggestion
def most_recent_job(job_type: str, session: Session = NEW_SESSION) ->
BaseJob | None:
```
##########
airflow/www/views.py:
##########
@@ -635,12 +635,14 @@ class AirflowBaseView(BaseView):
def render_template(self, *args, **kwargs):
# Add triggerer_job only if we need it
- if TriggererJob.is_needed():
- kwargs["triggerer_job"] =
lazy_object_proxy.Proxy(TriggererJob.most_recent_job)
+ if TriggererJobRunner.is_needed():
+ kwargs["triggerer_job"] = lazy_object_proxy.Proxy(
+ lambda: most_recent_job(TriggererJobRunner.job_type)
+ )
Review Comment:
```suggestion
kwargs["triggerer_job"] =
lazy_object_proxy.Proxy(TriggererJobRunner.most_revent_job)
```
Can't this use the helper method still?
##########
airflow/jobs/backfill_job.py:
##########
@@ -835,7 +843,7 @@ def _execute(self, session: Session = NEW_SESSION) -> None:
if not dagrun_infos:
if not self.run_at_least_once:
self.log.info("No run dates were found for the given dates and
dag interval.")
- return
+ return None
Review Comment:
(Nit: I'd rather keep the bare returns)
##########
airflow/task/task_runner/__init__.py:
##########
@@ -38,11 +39,11 @@
}
-def get_task_runner(local_task_job) -> BaseTaskRunner:
+def get_task_runner(local_task_job_runner: LocalTaskJobRunner) ->
BaseTaskRunner:
"""
- Get the task runner that can be used to run the given job.
+ Get the task runner that can be used to run with the give job runner.
Review Comment:
```suggestion
Get the task runner that can be used to run with the given job runner.
```
##########
airflow/www/views.py:
##########
@@ -660,7 +662,7 @@ def health(self):
scheduler_status = "unhealthy"
payload["metadatabase"] = {"status": "healthy"}
try:
- scheduler_job = SchedulerJob.most_recent_job()
+ scheduler_job = most_recent_job(SchedulerJobRunner.job_type)
Review Comment:
```suggestion
scheduler_job = SchedulerJob.most_recent_job()
```
And here too?
##########
airflow/jobs/backfill_job.py:
##########
@@ -921,6 +929,7 @@ def _execute(self, session: Session = NEW_SESSION) -> None:
executor.end()
self.log.info("Backfill done for DAG %s. Exiting.", self.dag)
+ return None
Review Comment:
(related nit: I'd rather we don't start the precedent of having these
returns at the end of functions)
--
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]
