Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 1 addition & 9 deletions airflow/cli/commands/task_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,6 @@
from airflow.utils import cli as cli_utils
from airflow.utils.cli import (
get_dag,
get_dag_by_deserialization,
get_dag_by_file_location,
get_dag_by_pickle,
get_dags,
Expand Down Expand Up @@ -364,14 +363,7 @@ def task_run(args, dag=None):
print(f'Loading pickle id: {args.pickle}')
dag = get_dag_by_pickle(args.pickle)
elif not dag:
if args.local:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hi @ephraimbuddy could you please give more context about why this should be removed?

it looks like your PR only tried to optimize when executing a task via fork. there is also an important part https://github.com/apache/airflow/pull/26750/files#diff-8b336b0706abbada03e8dce519e122384da2c913cc79ee4d638d0427a1342d41L47-L48

i think we should need dag = get_dag_by_deserialization(args.dag_id) please let me know your thoughts

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hi @ephraimbuddy could you please give more context about why this should be removed?

it looks like your PR only tried to optimize when executing a task via fork. there is also an important part https://github.com/apache/airflow/pull/26750/files#diff-8b336b0706abbada03e8dce519e122384da2c913cc79ee4d638d0427a1342d41L47-L48

i think we should need dag = get_dag_by_deserialization(args.dag_id) please let me know your thoughts

We noticed a significant increase in task run duration with the previous change. Using get_dag_by_deserialization only delays the dag parsing(for _start_by_fork). Looking at exec I think we should devise another way to get the serialized dag for use because using the previous change, we had to parse the dag in TaskRunner process for fork which causes delays.

try:
dag = get_dag_by_deserialization(args.dag_id)
except AirflowException:
print(f'DAG {args.dag_id} does not exist in the database, trying to parse the dag_file')
dag = get_dag(args.subdir, args.dag_id)
else:
dag = get_dag(args.subdir, args.dag_id)
dag = get_dag(args.subdir, args.dag_id, include_examples=False)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

include_examples is retrieved by default from configuration in get_dag. Can this change to have it always False be removed so that users who don't want examples to load can do so by using configuration? I was having a test case in my development PR that worked fine running an example dag and now it fails after rebase. This will be an issue to users too who might want to explore tasks command in cli with examples that worked before 2.4.0.

@ephraimbuddy ephraimbuddy Oct 12, 2022

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for this late reply. Can you share the test case?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was trying out a patch for this issue : #26555 . I just added an example of dynamic task mapping dag in this PR as a fixture : #26678 . I used to use this dag in my test case calling airflow tasks render which used to pass earlier but failed after rebase. I don't have the PR created yet.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
index 982aa31fd5..ab05bb56a9 100644
--- a/airflow/cli/commands/task_command.py
+++ b/airflow/cli/commands/task_command.py
@@ -578,13 +578,13 @@ def task_render(args):
         task, args.map_index, exec_date_or_run_id=args.execution_date_or_run_id, create_if_necessary="memory"
     )
     ti.render_templates()
-    for attr in task.__class__.template_fields:
+    for attr in task.template_fields:
         print(
             textwrap.dedent(
                 f"""        # ----------------------------------------------------------
         # property: {attr}
         # ----------------------------------------------------------
-        {getattr(task, attr)}
+        {getattr(ti.task, attr)}
         """
             )
         )
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
index c565f601d7..f3673fb2dd 100644
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -388,6 +388,54 @@ class TestCliTasks:
         assert 'echo "2016-01-01"' in output
         assert 'echo "2016-01-08"' in output
 
+    def test_mapped_task_render(self):
+        """
+        tasks render should render and displays templated fields for a given mapped task
+        """
+        with redirect_stdout(io.StringIO()) as stdout:
+            task_command.task_render(
+                self.parser.parse_args(
+                    [
+                        "tasks",
+                        "render",
+                        "test_mapped_classic",
+                        "consumer_literal",
+                        "2022-01-01",
+                        "--map-index",
+                        "0",
+                    ]
+                )
+            )
+        # the dag test_mapped_classic has op_args=[[1], [2], [3]], so the first mapped task should have op_args=[1]
+        output = stdout.getvalue()
+        assert "[1]" in output
+        assert "[2]" not in output
+        assert "[3]" not in output
+        assert "property: op_args" in output
+
+    def test_mapped_task_render_with_template(self):
+        """
+        tasks render should render and displays templated fields for a given mapped task
+        """
+        with redirect_stdout(io.StringIO()) as stdout:
+            task_command.task_render(
+                self.parser.parse_args(
+                    [
+                        "tasks",
+                        "render",
+                        "test_mapped_task_with_template",
+                        "some_command",
+                        "2022-01-01",
+                        "--map-index",
+                        "0",
+                    ]
+                )
+            )
+        # the dag test_mapped_classic has op_args=[[1], [2], [3]], so the first mapped task should have op_args=[1]
+        output = stdout.getvalue()
+        assert 'echo "2022-01-01"' in output
+        assert 'echo "2022-01-08"' in output
+
     def test_cli_run_when_pickle_and_dag_cli_method_selected(self):
         """
         tasks run should return an AirflowException when invalid pickle_id is passed
diff --git a/tests/dags/test_mapped_task_with_templates.py b/tests/dags/test_mapped_task_with_templates.py
new file mode 100644
index 0000000000..fe0b8d0c49
--- /dev/null
+++ b/tests/dags/test_mapped_task_with_templates.py
@@ -0,0 +1,36 @@
+# 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
+
+import datetime
+from textwrap import dedent
+
+from airflow import DAG
+from airflow.operators.bash import BashOperator
+
+with DAG(dag_id="test_mapped_task_with_template", start_date=datetime.datetime(2022, 1, 1)) as dag:
+    templated_command = dedent(
+        """
+    {% for i in range(5) %}
+        echo "{{ ds }}"
+        echo "{{ macros.ds_add(ds, 7)}}"
+    {% endfor %}
+    """
+    )
+    commands = [templated_command, "echo 1"]
+
+    BashOperator.partial(task_id="some_command").expand(bash_command=commands)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @ephraimbuddy for the patch. I will use it and also add you as co-author.

else:
# Use DAG from parameter
pass
Expand Down
8 changes: 3 additions & 5 deletions airflow/task/task_runner/standard_task_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ class StandardTaskRunner(BaseTaskRunner):
def __init__(self, local_task_job):
super().__init__(local_task_job)
self._rc = None
self.dag = local_task_job.task_instance.task.dag

def start(self):
if CAN_FORK and not self.run_as_user:
Expand Down Expand Up @@ -64,7 +65,6 @@ def _start_by_fork(self):
from airflow import settings
from airflow.cli.cli_parser import get_parser
from airflow.sentry import Sentry
from airflow.utils.cli import get_dag

# Force a new SQLAlchemy session. We can't share open DB handles
# between process. The cli code will re-create this as part of its
Expand Down Expand Up @@ -92,10 +92,8 @@ def _start_by_fork(self):
dag_id=self._task_instance.dag_id,
task_id=self._task_instance.task_id,
):
# parse dag file since `airflow tasks run --local` does not parse dag file
dag = get_dag(args.subdir, args.dag_id)
args.func(args, dag=dag)
return_code = 0
args.func(args, dag=self.dag)
return_code = 0
except Exception as exc:
return_code = 1

Expand Down
19 changes: 6 additions & 13 deletions airflow/utils/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
from typing import TYPE_CHECKING, Callable, TypeVar, cast

from airflow import settings
from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.utils import cli_action_loggers
from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler
Expand Down Expand Up @@ -205,7 +206,9 @@ def _search_for_dag_file(val: str | None) -> str | None:
return None


def get_dag(subdir: str | None, dag_id: str) -> DAG:
def get_dag(
subdir: str | None, dag_id: str, include_examples=conf.getboolean('core', 'LOAD_EXAMPLES')
) -> DAG:
"""
Returns DAG of a given dag_id

Expand All @@ -216,28 +219,18 @@ def get_dag(subdir: str | None, dag_id: str) -> DAG:
from airflow.models import DagBag

first_path = process_subdir(subdir)
dagbag = DagBag(first_path)
dagbag = DagBag(first_path, include_examples=include_examples)
if dag_id not in dagbag.dags:
fallback_path = _search_for_dag_file(subdir) or settings.DAGS_FOLDER
logger.warning("Dag %r not found in path %s; trying path %s", dag_id, first_path, fallback_path)
dagbag = DagBag(dag_folder=fallback_path)
dagbag = DagBag(dag_folder=fallback_path, include_examples=include_examples)
if dag_id not in dagbag.dags:
raise AirflowException(
f"Dag {dag_id!r} could not be found; either it does not exist or it failed to parse."
)
return dagbag.dags[dag_id]


def get_dag_by_deserialization(dag_id: str) -> DAG:
from airflow.models.serialized_dag import SerializedDagModel

dag_model = SerializedDagModel.get(dag_id)
if dag_model is None:
raise AirflowException(f"Serialized DAG: {dag_id} could not be found")

return dag_model.dag


def get_dags(subdir: str | None, dag_id: str, use_regex: bool = False):
"""Returns DAG(s) matching a given regex or dag_id"""
from airflow.models import DagBag
Expand Down
64 changes: 0 additions & 64 deletions tests/cli/commands/test_task_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,38 +159,6 @@ def test_test_filters_secrets(self, capsys):
task_command.task_test(args)
assert capsys.readouterr().out.endswith(f"{not_password}\n")

@mock.patch("airflow.cli.commands.task_command.get_dag_by_deserialization")
@mock.patch("airflow.cli.commands.task_command.LocalTaskJob")
def test_run_get_serialized_dag(self, mock_local_job, mock_get_dag_by_deserialization):
"""
Test using serialized dag for local task_run
"""
task_id = self.dag.task_ids[0]
args = [
'tasks',
'run',
'--ignore-all-dependencies',
'--local',
self.dag_id,
task_id,
self.run_id,
]
mock_get_dag_by_deserialization.return_value = SerializedDagModel.get(self.dag_id).dag

task_command.task_run(self.parser.parse_args(args))
mock_local_job.assert_called_once_with(
task_instance=mock.ANY,
mark_success=False,
ignore_all_deps=True,
ignore_depends_on_past=False,
ignore_task_deps=False,
ignore_ti_state=False,
pickle_id=None,
pool=None,
external_executor_id=None,
)
mock_get_dag_by_deserialization.assert_called_once_with(self.dag_id)

def test_cli_test_different_path(self, session):
"""
When thedag processor has a different dags folder
Expand Down Expand Up @@ -265,38 +233,6 @@ def test_cli_test_different_path(self, session):
# verify that the file was in different location when run
assert ti.xcom_pull(ti.task_id) == new_file_path.as_posix()

@mock.patch("airflow.cli.commands.task_command.get_dag_by_deserialization")
@mock.patch("airflow.cli.commands.task_command.LocalTaskJob")
def test_run_get_serialized_dag_fallback(self, mock_local_job, mock_get_dag_by_deserialization):
"""
Fallback to parse dag_file when serialized dag does not exist in the db
"""
task_id = self.dag.task_ids[0]
args = [
'tasks',
'run',
'--ignore-all-dependencies',
'--local',
self.dag_id,
task_id,
self.run_id,
]
mock_get_dag_by_deserialization.side_effect = mock.Mock(side_effect=AirflowException('Not found'))

task_command.task_run(self.parser.parse_args(args))
mock_local_job.assert_called_once_with(
task_instance=mock.ANY,
mark_success=False,
ignore_all_deps=True,
ignore_depends_on_past=False,
ignore_task_deps=False,
ignore_ti_state=False,
pickle_id=None,
pool=None,
external_executor_id=None,
)
mock_get_dag_by_deserialization.assert_called_once_with(self.dag_id)

@mock.patch("airflow.cli.commands.task_command.LocalTaskJob")
def test_run_with_existing_dag_run_id(self, mock_local_job):
"""
Expand Down