Add Airflow variable used to configure overrides for task timeouts #976

stacimc · 2023-01-28T00:01:42Z

Fixes

Fixes WordPress/openverse#1437 by @AetherUnbound

Description

Adds an Airflow Variable that can be used optionally to temporarily alter the execution timeouts for any task in a provider DAG, using a task_id_pattern. The timeout overrides are configured in the format <days>d:<hours>h:<minutes>m:<seconds>s. Example using the CONFIGURATION_OVERRIDES Airflow variable:

{
  # Example showing overriding timeouts for multiple tasks within a DAG
  "brooklyn_museum_workflow": [
    {
      "task_id_pattern": "pull_image_data",
      "timeout": "10h"  # 10 hrs
    },
    {
      "task_id_pattern": "report_load_completion",
      "timeout": "4s"  # 4 seconds
    }
  ],
  # Example showing task_id_pattern matching reingestion tasks (that have `day_shift_x` suffixes)
  "metropolitan_museum_reingestion_workflow": [
    {
      "task_id_pattern": "upsert_data",
      "timeout": "6d:11h"  # 6 days, 11 hours
    }
  ],
  "inaturalist_workflow": [
    {
      "task_id_pattern": "load_taxa",
      "timeout": "6d:10h:9s"  # 6 days, 10 hours, 9 seconds
    },
    # Example of a mapped task
    {
      "task_id_pattern": "load_transformed_data",
      "timeout": "10m"  # 10 minutes
    }
  ]
}

This is intentionally implemented in such a way that it can be extended very quickly to override other properties of tasks. I had retries in mind but did not include it in this PR for simplicity.

Testing Instructions

You can use the example CONFIGURATION_OVERRIDES provided above, and play around with adding other configurations. For each timeout you configure, run the DAG and check the Task Instance Details page for the task to verify that the execution_timeout is set to the one you configured. You can try setting really short timeouts on long-running tasks like pull_data to verify that the timeouts will actually be hit.

Make sure to test:

tasks on a regular ingestion workflow, such as pull_image_data
tasks on a reingestion workflow with day_shift applied. Example, try upsert_data and verify that the timeout gets applied to all upsert_data_day_shift_x tasks within the workflow
tasks on iNaturalist (our only example of a 'custom' workflow)
- specifically, iNaturalist's load_transformed_data task, which is a Mapped task. The timeout should be applied to each mapped task

Notes:

Improperly formatted overrides are simply ignored, rather than raising an error which will break all provider DAGs.

Checklist

My pull request has a descriptive title (not a vague title like
Update index.md).
My pull request targets the default branch of the repository (main) or
a parent feature branch.
My commit messages follow best practices.
My code follows the established code style of the repository.
I added or updated tests for the changes I made (if applicable).
I added or updated documentation (if applicable).
I tried running the project locally and verified that there are no visible
errors.
I ran the DAG documentation generator (if applicable).

Developer Certificate of Origin

Developer Certificate of Origin
Version 1.1

Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
1 Letterman Drive
Suite D4700
San Francisco, CA, 94129

Everyone is permitted to copy and distribute verbatim copies of this
license document, but changing it is not allowed.


Developer's Certificate of Origin 1.1

By making a contribution to this project, I certify that:

(a) The contribution was created in whole or in part by me and I
    have the right to submit it under the open source license
    indicated in the file; or

(b) The contribution is based upon previous work that, to the best
    of my knowledge, is covered under an appropriate open source
    license and I have the right under that license to submit that
    work with modifications, whether created in whole or in part
    by me, under the same open source license (unless I am
    permitted to submit under a different license), as indicated
    in the file; or

(c) The contribution was provided directly to me by some other
    person who certified (a), (b) or (c) and I have not modified
    it.

(d) I understand and agree that this project and the contribution
    are public and that a record of the contribution (including all
    personal information I submit with it, including my sign-off) is
    maintained indefinitely and may be redistributed consistent with
    this project or the open source license(s) involved.

AetherUnbound · 2023-02-14T18:25:18Z

Dang, this is really tough! The flexibility of defining agnostic, per-task timeouts is neat. But I think about what I had to do for WordPress/openverse#1301 - because of the reingestion workflows, we'd probably also need to make this a task ID pattern for matching timeouts 😕 How much additional complexity would that add do you think?

I think it's fair to try and target the typical ProviderDataIngester classes with this, and maybe try to fit iNaturalist into that mold as you're suggesting. Even that utility alone would be beneficial!

… overrides

stacimc · 2023-02-22T20:14:16Z

@AetherUnbound It ended up being easier than I thought! Updated version of the PR now allows arbitrarily overriding timeout on any task in a provider DAG, including iNaturalist.

AetherUnbound

Woohoo! What an excellent feature this will unlock for us. I have one request regarding the timeout format that I'd like your thoughts on 🕐

openverse_catalog/dags/providers/provider_workflows.py

AetherUnbound · 2023-02-28T01:28:16Z

openverse_catalog/dags/providers/provider_reingestion_workflows.py

-            # Override the dag_id
-            self.dag_id = f"{self.provider_name}_reingestion_workflow"
-            return
+            _, provider_name = self._get_module_info()


This is a much better approach! 😄

AetherUnbound · 2023-02-28T01:34:43Z

openverse_catalog/dags/providers/provider_dag_factory.py

+            # If the task is a MappedOperator, apply the timeout to partial
+            # kwargs to ensure it is set on each mapped task.
+            if isinstance(task, MappedOperator):
+                task.partial_kwargs["execution_timeout"] = timeout_override
+            else:
+                task.execution_timeout = timeout_override


Wow, what a find! 😱

AetherUnbound

This is an excellent change! I know it can be hard to test, but I love the ability to just have the time range we need specified (e.g. 10h). Very excited for this!

AetherUnbound · 2023-03-02T23:29:30Z

openverse_catalog/dags/providers/provider_workflows.py

+        match unit:
+            case "d":
+                days = int(count)
+            case "h":
+                hours = int(count)
+            case "m":
+                minutes = int(count)
+            case "s":
+                seconds = int(count)


A case statement!! Wow, my first ever! 😄

openverse_catalog/dags/providers/provider_workflows.py

krysal

It's so flexible! Awesome, it worked like a charm! 💯

Add Airflow variable used to configure overrides for task timeouts

2c242a3

stacimc added 🟩 priority: low Low priority and doesn't need to be rushed ✨ goal: improvement Improvement to an existing user-facing feature 💻 aspect: code Concerns the software code in the repository labels Jan 28, 2023

stacimc self-assigned this Jan 28, 2023

stacimc added 6 commits February 17, 2023 17:33

Don't raise error when time is incorrectly formatted

4cdadd1

Lint

db8a01f

Ensure dag_id is set for reingestion flows before attempting to apply…

73f5119

… overrides

Allow overriding any task timeout by pattern

97ddb22

Add support for mapped tasks

2765427

Lint docstring

ce5661a

stacimc marked this pull request as ready for review February 22, 2023 20:14

stacimc requested a review from a team as a code owner February 22, 2023 20:14

stacimc requested review from krysal and AetherUnbound February 22, 2023 20:14

AetherUnbound reviewed Feb 28, 2023

View reviewed changes

Update time parsing

ad350cb

AetherUnbound approved these changes Mar 2, 2023

View reviewed changes

stacimc mentioned this pull request Mar 7, 2023

Add a DAG for backfilling license_url when meta_data is null #1005

Merged

8 tasks

krysal approved these changes Mar 8, 2023

View reviewed changes

obulat added 🧱 stack: catalog Related to the catalog and Airflow DAGs and removed 🧱 stack: catalog Related to the catalog and Airflow DAGs labels Mar 9, 2023

stacimc merged commit 63cfe8a into main Mar 9, 2023

stacimc deleted the add/configuration-overrides-for-dag-timeouts branch March 9, 2023 18:24

stacimc mentioned this pull request Mar 23, 2023

Log last query_params hit before AirflowTaskTimeout #1058

Merged

7 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Airflow variable used to configure overrides for task timeouts #976

Add Airflow variable used to configure overrides for task timeouts #976

stacimc commented Jan 28, 2023 •

edited

Loading

AetherUnbound commented Feb 14, 2023

stacimc commented Feb 22, 2023

AetherUnbound left a comment

AetherUnbound Feb 28, 2023

AetherUnbound Feb 28, 2023

AetherUnbound left a comment

AetherUnbound Mar 2, 2023

krysal left a comment

Add Airflow variable used to configure overrides for task timeouts #976

Add Airflow variable used to configure overrides for task timeouts #976

Conversation

stacimc commented Jan 28, 2023 • edited Loading

Fixes

Description

Testing Instructions

Checklist

Developer Certificate of Origin

AetherUnbound commented Feb 14, 2023

stacimc commented Feb 22, 2023

AetherUnbound left a comment

Choose a reason for hiding this comment

AetherUnbound Feb 28, 2023

Choose a reason for hiding this comment

AetherUnbound Feb 28, 2023

Choose a reason for hiding this comment

AetherUnbound left a comment

Choose a reason for hiding this comment

AetherUnbound Mar 2, 2023

Choose a reason for hiding this comment

krysal left a comment

Choose a reason for hiding this comment

stacimc commented Jan 28, 2023 •

edited

Loading