Fix: Add task context labels to driver and executor pods for SparkKubernetesOperator reattach_on_restart functionality (#41211)

[asb108]

Closes: #41211 where the SparkKubernetesOperator's reattach_on_restart functionality doesn't work correctly.

Problem

When reattach_on_restart is enabled, the SparkKubernetesOperator tries to find the driver pod by looking for pods with specific task context labels (dag_id, task_id, run_id). However, these labels are not actually added to the driver and executor pods when creating them, causing the reattach functionality to fail.

Solution

This PR adds code to the execute method of the SparkKubernetesOperator class to add task context labels to both the driver and executor pods when reattach_on_restart is enabled. This allows the operator to find the existing driver pod if the scheduler restarts.

Testing

I've performed comprehensive testing to ensure the fix works correctly:

1. Unit tests that verify the task context labels are correctly added when reattach_on_restart is enabled
2. Tests that verify the default behavior remains unchanged when reattach_on_restart is disabled
3. Integration tests in a real Kubernetes environment using Kind
4. Tests with different Spark application configurations to ensure compatibility

All tests passed, confirming that the fix works as expected.

[boring-cyborg]

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:

• Pay attention to the quality of your code (ruff, mypy and type annotations). Our pre-commits will help you with that.
• In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
• Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
• Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
• Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
• Be sure to read the Airflow Coding style.
• Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
  Apache Airflow is a community-driven project and together we are making it better 🚀.
  In case of doubts contact the developers at:
  Mailing List: dev@airflow.apache.org
  Slack: https://s.apache.org/airflow-slack

[asb108]

Hi @hussein-awala @jedcunningham @eladkal,

I've pushed the changes to address the raised issue and the previously failing test. When you have a moment, could you please review the updates? I'm new to the Airflow repository, and this is my first PR, so I really appreciate any feedback or suggestions you might have.

Thanks for your time and support!

[asb108]

Thanks for triggering tests @eladkal,
I'm seeing consistent failures in the Airflow providers test suite, in the job runs at https://github.com/apache/airflow/actions/runs/15146342927/job/42604235636?pr=50803 (ref: 6fcc9f1). The jobs are being killed with exit code 137, indicating an out-of-memory (OOM) issue. There are no other explicit errors—detailed test output is missing as the process is terminated before it can be printed.

Should I consider optimizing the tests for memory usage? I'd appreciate any suggestions on improving test reliability or adjusting the job configuration to reduce resource consumption.

Thank you

[potiuk]

Should I consider optimizing the tests for memory usage?

Yes - i think it's best you track down where it comes from. It's hard to say and we are not experts in all integrations, but yes indeed it looks like somewhere memory is leaked - but where and how - hard to say. Note that there are usually some logs uploaded as artifacts (you can scroll down summary page) - maybe they are helpful. I also added "debug ci resources" label that might help you to see if indeed memory is growing and when more easily - it will print resource stats in the periodic dump of your test run - you can trace it when you rebase and re-run the PR.

[asb108]

Hi @eladkal ,
the conflicts are resolved.

Sorry missed that. Can tou rebase again?

Okay

[boring-cyborg]

Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.