Page MenuHomePhabricator

Spike: figure out how to efficiently send XComs to Airflow dynamically mapped tasks
Open, Needs TriagePublic

Description

For our DAGs with bigger dynamically mapped tasks ( refine_to_hive_hourly and mw_content_xml_export_*_monthly ) we have hit multiple issues to be able to send in XComs in.

The current solution is to have an upstream task that generates the XComs, and to access them like so in the dynamically mapped task:

def get_spark_tuning_from_xcom_as_str(
    task_id: str,
    config_name: str,
    default: str = "",
) -> str:
    """Get a specific configuration from the XComs of assign_spark_tunings."""
    return (
        f"{{{{ ti.xcom_pull(task_ids='{task_id}')[ti.map_index].get('{config_name}', '{default}') }}}}"  # noqa
    )

This works, but presumably puts significant pressure on the DB, as each call to get_spark_tuning_from_xcom_as_str() would resolve the whole XCom array, to then just get a particular dict, and from that dict extract just one element. And we hit this method N times per dynamic task, as in this example filling up a Spark conf:

conf={
    "spark.jars": props.refinery_job_jar_path,
    "spark.sql.shuffle.partitions": get_spark_tuning_from_xcom_as_str(
        "small_wikis_tunings", "spark.sql.shuffle.partitions", "1024"
    ),
    "spark.sql.broadcastTimeout": get_spark_tuning_from_xcom_as_str(
        "small_wikis_tunings", "spark.sql.broadcastTimeout", "300"
    ),
    "spark.reducer.maxReqsInFlight": get_spark_tuning_from_xcom_as_str(
        "small_wikis_tunings", "spark.reducer.maxReqsInFlight", "2147483647"
    ),
    "spark.shuffle.io.maxRetries": get_spark_tuning_from_xcom_as_str(
        "small_wikis_tunings", "spark.shuffle.io.maxRetries", "3"
    ),
    "spark.network.timeout": get_spark_tuning_from_xcom_as_str(
        "small_wikis_tunings", "spark.network.timeout", "120"
    ),
    "spark.task.maxFailures": get_spark_tuning_from_xcom_as_str(
        "small_wikis_tunings", "spark.task.maxFailures", "4"
    ),
},

In this spike we want to explore:

  • the actual impact to the DB. Do we have metrics on this?
  • If we find significant DB stress, find if the there is a way to reduce the DB callls
  • report our findings here, perhaps on a PoC DAG that can be used to refactor other tasks

context:
https://gitlab.wikimedia.org/repos/data-engineering/airflow-dags/-/merge_requests/1652#note_167415
and
https://gitlab.wikimedia.org/repos/data-engineering/airflow-dags/-/merge_requests/1652#note_167274