justinpakzad opened a new pull request, #68259:
URL: https://github.com/apache/airflow/pull/68259
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
**Description**
Snowpark Container Services lets users run containerized workloads directly
within their Snowflake account. Things like ML training jobs, data processing
pipelines, or any custom application packaged as a Docker image. Jobs are
submitted via EXECUTE JOB SERVICE, which runs containers on a compute pool and
terminates once the code exits. Currently there's no operator for orchestrating
these jobs. Users would have to fall back to raw SQL via
SQLExecuteQueryOperator and manually handle polling, log retrieval, error
handling, and service teardown.
This PR adds SnowparkContainerJobOperator which:
- Submits the job asynchronously via EXECUTE JOB SERVICE
- Polls DESCRIBE SERVICE until the job reaches a terminal state (DONE,
FAILED, etc.)
- Fetches container logs via SYSTEM$GET_SERVICE_LOGS for all replicas and
surfaces them in the Airflow task logs
- Supports both stage-based specs (spec + spec_stage) and inline YAML
(spec_text)
- Supports fire-and-forget mode (wait_for_completion=False) for
long-running jobs as well as other optional parameters (drop_on_completion,
replicas, query_warehouse, etc.)
**Testing:**
- Unit tests covering parameter validation, SQL generation, polling logic,
log retrieval, error handling for unexpected statuses, on_kill cleanup, and
execution paths.
- Manually tested end-to-end against a real Snowflake account with an active
compute pool, image repository, and pushed container image verified the success
path (job completes, logs retrieved) and failure scenarios (job fails, error
logs surfaced, operator raises appropriately).
Added an example DAG and operator documentation. Running the example DAG
requires SPCS infrastructure to be provisioned beforehand (compute pool, image
repository, pushed container image, etc.) those steps are covered in the
operator's how-to guide.
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [x] Yes (Claude Opus 4.6)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request
Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)**
for more information. Note: commit author/co-author name and email in commits
become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal
([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals))
is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License
Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment:
`{pr_number}.significant.rst`, in
[airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
You can add this file in a follow-up commit after the PR is created so you
know the PR number.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
