apache / apache/airflow

Jinja templating capability is not working for the parameter (additional args) in spark job submit for DatabricksWorkflowTaskGroup

Open
#46,563 12 comments 4 reactions 0 assignees View on GitHub
area:providers good first issue kind:bug provider:databricks
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 10h
Merged PRs (30d)
483

Description

### Apache Airflow version

2.10.4

### If "Other Airflow 2 version" selected, which one?

_No response_

### What happened?

While creating workflow in Databricks using `DatabricksWorkflowTaskGroup`, I have a situation to pass some parameters to spark submit which are passed as jinja template.

Sample task creation code with DatabricksWorkflowTaskGroup,
```
databricks_workflow = DatabricksWorkflowTaskGroup(
group_id="test_wf",
databricks_conn_id=DATABRICKS_CONNECTION_ID,
add_suffix_on_collision=True,
job_clusters=job_clusters_spec,
)

with databricks_workflow:
task1 = build_task_1(dag)
task2 = build_task_2(dag)
task1 >> task2
```

From the above sample workflow creation code, `build_task_1` and `build_task_2` returns task which is my custom build operator on extending `DatabricksTaskBaseOperator`.

Sample code for custom build operator,

```
class CustomDatabricksWorkflowTaskOperator(DatabricksTaskBaseOperator):
"""Custom operator wrapping DatabricksTaskOperator"""

template_fields = (
"application_args",
"spark_conf",
"driver_jvm_props",
)

def __init__(
self,
*args,
class_name: str = "",
spark_conf: dict[str, str] | str = None,
application_args: List[str] | str | None = None,
driver_jvm_props: dict[str, str] | None = None,
application_conf_file: str = "test.conf",
libraries: List[dict[str, str]] = None,
**kwargs,
):
self.class_name = class_name
self.application_conf_file = application_conf_file
self.library = libraries if libraries is not None else _get_default_libraries()
self.application_args = application_args
super().__init__(*args, **kwargs)

self.driver_jvm_props = driver_jvm_props
self.spark_conf = spark_cone

def _get_task_base_json(self) -> dict[str, Any]:
return {
"run_if": self.databricks_trigger_rule,
'spark_jar_task': {
'main_class_name': self.class_name,
parameters": [
self.application_args,
"{{ get_test_template_value(ti,'test_key') }}" // I have hard-coded here for ref. In ideal case, this comes as part of self.application_args which has been added in template_fields
]
},
'libraries': self.library,
}

def execute(self, context):
super().execute(context)
```

From the above code sample, `{{ get_test_template_value(ti,'test_key') }}` is not getting resolved at the time of `launch` task from DatabricksWorkflowTaskGroup. But the same value got resolved at the time of task execution in airflow.

Sample json from `launch` task:
```
"tasks": [
{
"task_key": "task1",
"depends_on": [],
"run_if": "NONE_FAILED",
"spark_jar_task": {
"main_class_name": "com.sample.task1",
"parameters": [
"app.conf",
"{{ get_test_template_value(ti,'test_key') }}"
]
},
"libraries": [
{
"jar": "test.jar"
}
],
"job_cluster_key": "job_cluster"
},
```
Note: `get_test_template_value` function is added as a macro while DAG creation.

Though I have included the `application_args` as `template_fields` section in the CustomDatabricksWorkflowTaskOperator, it is unable to resolve the template.

### What you think should happen instead?

Expectation is to jinja template resolution should happen while the launch tasks trigger API requests to Databricks to create the workflow. But since the template resolution is happening lazily here i.e., at the time of task execution in airflow which is late in this case as the respective task is already created in Databricks at the time of `launch` and Databricks doesn't have to way to get the template resolved value.

### How to reproduce

1. Create a DAG with `DatabricksWorkflowTaskGroup`
2. Add couple of tasks
3. Create a custom Databricks Task operator by extending `DatabricksTaskBaseOperator`
4. Add one of the constructor parameter in `template_fields` section which needs to be tested
5. Override `_get_task_base_json` method to return a templated field/value as of the return dictionary
6. Trigger DAG
7. In launch task, you can notice from logs that the parameter property will pass the string as-is without resolving the actual value.
8. In the actual task execution, the actual value can be seen in the rendered template section.

### Operating System

PRETTY_NAME="Debian GNU/Linux 12 (bookworm)", NAME="Debian GNU/Linux"

### Versions of Apache Airflow Providers

NA

### Deployment

Other Docker-based deployment

### Deployment details

We are building our own docker image with base image as `apache/airflow:2.10.4-python3.10` and use the same for deployment.

### Anything else?

Other approaches which I tried inside `_get_task_base_json` method of the custom operation which is responsible for returning the task JSON.

1. Tried to force the render by using `super().render_template(self.application_args, context)` and `super().render_template_fields(context)` which unfortunately was not possible as the context was not part of the `_get_task_base_json` method scope.
2. As this is an abstract method and I had to override, unable to pass context object to the step1.

Due to the above constraint, render templating as well as pulling from xcom also is not supported. (In other words, all the capabilities which we get from using context object is not possible).

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [x] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)

Contributor guide

Open the contributing guide

Research direction

Start by tracing DatabricksWorkflowTaskGroup's launch task and the DatabricksTaskBaseOperator implementation, especially the overridden _get_task_base_json method described in the issue. Follow how template context reaches task execution and how the launch API payload is assembled; done means templated parameters are resolved before the Databricks workflow is created, with coverage for the reproduced custom-operator case.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.