Swap Dag Parsing to use the TaskSDK machinery. #44972

ashb · 2024-12-16T23:25:04Z

As part of Airflow 3 DAG definition files will have to use the Task SDK for
all their classes, and anything involving running user code will need to be
de-coupled from the database in the user-code process.

This change moves all of the "serialization" change up to the
DagFileProcessorManager, using the new function introduced in #44898 and the
"subprocess" machinery introduced in #44874.

Important Note: this change does not remove the ability for dag processes
to access the DB for Variables etc. That will come in a future change.

Some key parts of this change:

It builds upon the WatchedSubprocess from the TaskSDK. Right now this puts a
nasty/unwanted depenednecy between the Dag Parsing code upon the TaskSDK.
This will be addressed before release (we have talked about introducing a
new "apache-airflow-base-executor" dist where this subprocess+supervisor
could live, as the "execution_time" folder in the Task SDK is more a feature
of the executor, not of the TaskSDK itself)
A number of classes that we need to send between processes have been
converted to Pydantic for ease of serialization.
In order to not have to serialize everything in the subprocess and deserialize everything
in the parent Manager process, we have created a LazyDeserializedDAG class
that provides lazy access to much of the properties needed to create update
the DAG related DB objects, without needing to fully deserialize the entire
DAG structure.
Classes switched to attrs based for less boilerplate in constructors.
Internal timers convert to time.monotonic where possible, and time.time
where not, we only need second diff between two points, not datetime objects
With the earlier removal of "sync mode" for SQLite in Remove "single process" restrictions on SQLite in favour of using WAL mode #44839 the need for
separate TERMIANTE and END messages over the control socket can go

Co-authored-by: Jed Cunningham [email protected]
Co-authored-by: Daniel Imberman [email protected]

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

ashb · 2024-12-16T23:25:47Z

The tests aren't 100% finished yet.

And this change is larger than I would have liked, but at least it's a net-negative change

tests/listeners/test_dag_import_error_listener.py

airflow/dag_processing/manager.py

airflow/dag_processing/processor.py

ashb · 2024-12-16T23:49:38Z

I don't expect tests to pass yet, but I want to give people the chance to see this PR, and I know @jedcunningham is waiting on this for some of his DAG versioning work.

kaxil · 2024-12-17T06:59:15Z

Right now this puts a
nasty/unwanted depenednecy between the Dag Parsing code upon the TaskSDK.

I think that is unavoidable, as user code will come from Task SDK

This will be addressed before release (we have talked about introducing a
new "apache-airflow-base-executor" dist where this subprocess+supervisor
could live, as the "execution_time" folder in the Task SDK is more a feature
of the executor, not of the TaskSDK itself)

Right, I think processor will have to depend on both Task SDK (user-facing code) + Base Executor dist -- after that separation

kaxil

Partial review, will be coming back to it in an hour

airflow/callbacks/callback_requests.py

airflow/dag_processing/collection.py

airflow/models/dagcode.py

airflow/dag_processing/collection.py

airflow/dag_processing/manager.py

airflow/dag_processing/collection.py

airflow/dag_processing/manager.py

airflow/callbacks/callback_requests.py

airflow/dag_processing/processor.py

airflow/serialization/serialized_objects.py

tests/dag_processing/test_manager.py

tests/dag_processing/test_collection.py

ashb · 2024-12-17T11:45:49Z

airflow/dag_processing/manager.py

+
+
+@attrs.define
+class DagFileProcessorManager:


Suggested change

class DagFileProcessorManager:

class FileParsingSupervisor:

and let's rename the file appropriately.

ashb · 2024-12-17T11:47:13Z

airflow/dag_processing/processor.py

-
-class DagFileProcessor(LoggingMixin):
+@attrs.define()
+class DagFileProcessorProcess(WatchedSubprocess):


Suggested change

class DagFileProcessorProcess(WatchedSubprocess):

class FileParserProcess(WatchedSubprocess):

Yeah sounds good

I'm going to do these rename as a separate PR right after this one -- the diff is big enough as it is.

airflow/dag_processing/processor.py

airflow/serialization/serialized_objects.py

task_sdk/src/airflow/sdk/log.py

amoghrajesh

I mainly reviewed the processor and manager files. The other changes seem mostly reactive. Overall, I like the improvements, especially the reuse of the execution time machinery here. Few initial comments, nothing serious but mostly nits.

amoghrajesh · 2024-12-18T08:28:58Z

airflow/dag_processing/processor.py

-
-class DagFileProcessor(LoggingMixin):
+@attrs.define()
+class DagFileProcessorProcess(WatchedSubprocess):


Yeah sounds good

airflow/dag_processing/processor.py

airflow/serialization/serialized_objects.py

airflow/dag_processing/manager.py

tests/dag_processing/test_processor.py

ashb · 2024-12-18T23:09:02Z

Right, I think this should now pass the tests, the only thing I'm not sure about this is the xfail I've put for the "simple ti roundtrip exec config tests" -- Either we should remove it or make it work, but I'm not sure if we need to pass down executor config via TI anymore

@kaxil Any ideas the best plan for that one?

ashb · 2024-12-18T23:09:26Z

(I still need to rename a class and file, but that is a non-meaningful/non-review-impacting change.

airflow/dag_processing/collection.py

airflow/serialization/serialized_objects.py

kaxil · 2024-12-19T08:43:53Z

Right, I think this should now pass the tests, the only thing I'm not sure about this is the xfail I've put for the "simple ti roundtrip exec config tests" -- Either we should remove it or make it work, but I'm not sure if we need to pass down executor config via TI anymore

@kaxil Any ideas the best plan for that one?

Since we are planning to handle callbacks via Executor/worker interface too -- don't think we need to pass it explicitly from TI, instead just handle it on the server side before sending TI/request to the worker.

kaxil

Some minor comments and the renaming of file & classes can happen in a separate PR too

airflow/dag_processing/manager.py

airflow/dag_processing/processor.py

ashb · 2024-12-19T12:01:34Z

I'm going to run this with full tests, I want the kube tests to see if there is something broken not covered by unit tests

As part of Airflow 3 DAG definition files will have to use the Task SDK for all their classes, and anything involving running user code will need to be de-coupled from the database in the user-code process. This change moves all of the "serialization" change up to the DagFileProcessorManager, using the new function introduced in apache#44898 and the "subprocess" machinery introduced in apache#44874. **Important Note**: this change does not remove the ability for dag processes to access the DB for Variables etc. That will come in a future change. Some key parts of this change: - It builds upon the WatchedSubprocess from the TaskSDK. Right now this puts a nasty/unwanted depenednecy between the Dag Parsing code upon the TaskSDK. This will be addressed before release (we have talked about introducing a new "apache-airflow-base-executor" dist where this subprocess+supervisor could live, as the "execution_time" folder in the Task SDK is more a feature of the executor, not of the TaskSDK itself.) - A number of classes that we need to send between processes have been converted to Pydantic for ease of serialization. - In order to not have to serialize everything in the subprocess and deserialize everything in the parent Manager process, we have created a `LazyDeserializedDAG` class that provides lazy access to much of the properties needed to create update the DAG related DB objects, without needing to fully deserialize the entire DAG structure. - Classes switched to attrs based for less boilerplate in constructors. - Internal timers convert to `time.monotonic` where possible, and `time.time` where not, we only need second diff between two points, not datetime objects. - With the earlier removal of "sync mode" for SQLite in apache#44839 the need for separate TERMIANTE and END messages over the control socket can go. Co-authored-by: Jed Cunningham <[email protected]> Co-authored-by: Daniel Imberman <[email protected]>

…e sense any more

ashb · 2024-12-19T12:18:19Z

tests/dag_processing/test_processor.py

+        assert "a.py" in resp.import_errors
+
+
+#     @conf_vars({("logging", "dag_processor_log_target"): "stdout"})


Leaving these comments out for now as I want to overhaul the logging in #45072

ashb added area:Scheduler including HA (high availability) scheduler area:task-execution-interface-aip72 AIP-72: Task Execution Interface (TEI) aka Task SDK area:task-sdk labels Dec 16, 2024

boring-cyborg bot added area:CLI area:serialization labels Dec 16, 2024

ashb force-pushed the dag-parsing-uses-task-sdk branch from 32176b9 to 1215213 Compare December 16, 2024 23:29

ashb commented Dec 16, 2024

View reviewed changes

tests/listeners/test_dag_import_error_listener.py Outdated Show resolved Hide resolved

ashb commented Dec 16, 2024

View reviewed changes

airflow/dag_processing/manager.py Show resolved Hide resolved

ashb commented Dec 16, 2024

View reviewed changes

airflow/dag_processing/manager.py Show resolved Hide resolved

ashb commented Dec 16, 2024

View reviewed changes

airflow/dag_processing/manager.py Show resolved Hide resolved

ashb commented Dec 16, 2024

View reviewed changes

airflow/dag_processing/processor.py Show resolved Hide resolved

ashb commented Dec 16, 2024

View reviewed changes

airflow/dag_processing/processor.py Show resolved Hide resolved

ashb commented Dec 16, 2024

View reviewed changes

airflow/dag_processing/processor.py Show resolved Hide resolved

ashb force-pushed the dag-parsing-uses-task-sdk branch from 1215213 to 9f04c14 Compare December 16, 2024 23:46