potiuk commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3952257954
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -203,7 +190,14 @@ def _submit_job(self) -> str:
def _poll_for_status(self) -> str:
"""Poll until the job reaches a terminal state."""
+ status = None
+ end_time = time.time() + self.timeout
Review Comment:
> - `time.monotonic()` for durations, not `time.time()`.
(`AGENTS.md` → Coding Standards)
This loop is entirely in-process, so `time.monotonic()` is the
standard-compliant choice and is immune to a wall-clock step (NTP correction,
manual clock change) silently shortening or extending the window. It'd be
`time.monotonic()` on both this line and the `time.time() >= end_time` check
below.
The trigger's `end_time` / `execution_deadline` are a different case and
correctly stay on `time.time()` — those are absolute deadlines that get
serialized and re-read in another process, where a monotonic value would be
meaningless.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -137,6 +117,8 @@ def __init__(
drop_on_completion: bool = True,
poll_interval: int = 10,
snowflake_conn_id: str = "snowflake_default",
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ timeout: int = 24 * 60 * 60,
Review Comment:
This default also changes the existing non-deferrable path:
`_poll_for_status` polled indefinitely in 6.15.0 / 6.16.1, and now a job still
running at 24h fails the task and — with the default `drop_on_completion=True`
— drops the service out from under it.
That's a sensible default, but it's a behaviour change for anyone already
running SPCS jobs longer than a day, and it isn't visible from the PR title.
Airflow's provider release managers regenerate `changelog.rst` from `git log`,
so a line in `providers/snowflake/docs/changelog.rst` (or at minimum a sentence
in the commit body) is the only place it'll get flagged for them.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -212,25 +206,41 @@ def _poll_for_status(self) -> str:
raise RuntimeError(f"Job {self.job_name} returned unexpected
status: {status}")
time.sleep(self.poll_interval)
- def _log_container_output(self, status: str) -> None:
- """Fetch and log container output for all replicas."""
+ def _log_container_output(self, status: str | None) -> None:
+ """Fetch and log container output for all replicas. Best-effort so it
never blocks cleanup."""
for instance_id in range(self.replicas):
sql = f"SELECT SYSTEM$GET_SERVICE_LOGS('{self.job_name}',
{instance_id}, '{self.container_name}')"
- response = self._run_one(sql)[0]
+ try:
+ response = self._run_one(sql)[0]
+ except Exception as e:
+ self.log.warning("Could not retrieve logs for instance_id %d:
%s", instance_id, e)
+ continue
if not response:
continue
if status != SnowparkContainerJobStatus.DONE:
self.log.error("Logs for instance_id %d:\n%s", instance_id,
response)
else:
self.log.info("Logs for instance_id %d:\n%s", instance_id,
response)
+ def _drop_service(self) -> None:
+ """Best-effort drop of the job service."""
+ try:
+ self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
+ except Exception as e:
+ self.log.error("Error dropping service %s: %s", self.job_name, e)
+
def on_kill(self) -> None:
"""Drop the running service on task kill."""
if self.job_name:
- try:
- self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
- except Exception as e:
- self.log.error("Error dropping service %s: %s", self.job_name,
e)
+ self._drop_service()
+
+ def _handle_final_status(self, status: str) -> None:
+ """Log container output, fail unless the job is DONE, and optionally
drop the service on success."""
+ self._log_container_output(status)
+ if status != SnowparkContainerJobStatus.DONE:
+ raise RuntimeError(f"Job '{self.job_name}' finished with status:
{status}")
+ if self.drop_on_completion:
+ self._drop_service()
Review Comment:
Small behaviour change worth being deliberate about. Before this PR the
success-path drop was:
```python
if self.drop_on_completion:
self._hook.run(f"DROP SERVICE IF EXISTS {self.job_name}")
```
— no `try`, so a drop that failed failed the task. Routing it through
`_drop_service` (line 225) now logs at `error` level and lets the task succeed,
which leaves a service running and billing with no failed task to alert on.
Swallowing may well be what you want here — the job did succeed, and failing
an otherwise-good task on a cleanup error is harsh. But `on_kill` was
best-effort because it runs during teardown, whereas this is the normal
completion path, and the failure is now silent. Either keeping the raise on the
completion path or a one-line note on why it's best-effort would make the
choice legible.
##########
providers/snowflake/src/airflow/providers/snowflake/triggers/snowpark_containers.py:
##########
@@ -0,0 +1,166 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+from airflow.providers.common.sql.hooks.handlers import fetch_one_handler
+from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.providers.snowflake.utils.snowpark_containers import (
+ NON_TERMINAL_STATUSES,
+ TERMINAL_STATUSES,
+)
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class SnowparkContainerJobTrigger(BaseTrigger):
+ """
+ Poll a Snowpark Container Services job until it reaches a terminal status.
+
+ :param job_name: name of the submitted job service to poll.
+ :param snowflake_conn_id: reference to the Snowflake connection id.
+ :param poll_interval: seconds to sleep between ``DESCRIBE SERVICE`` polls.
+ :param end_time: epoch deadline (``time.time()`` seconds) after which a
``timeout``
+ event is emitted.
+ :param execution_deadline: (Optional) absolute timestamp (in seconds since
the epoch) after
+ which the task is considered timed out. (default: None)
+ :param database: (Optional) name of database. (default: None)
+ :param schema: (Optional) name of schema. (default: None)
+ :param role: (Optional) name of role. (default: None)
+ :param warehouse: (Optional) name of warehouse. (default: None)
+ """
+
+ def __init__(
+ self,
+ job_name: str,
+ snowflake_conn_id: str,
+ poll_interval: float,
+ end_time: float,
+ execution_deadline: float | None = None,
+ database: str | None = None,
+ schema: str | None = None,
+ role: str | None = None,
+ warehouse: str | None = None,
+ ) -> None:
+ super().__init__()
+ self.job_name = job_name
+ self.snowflake_conn_id = snowflake_conn_id
+ self.poll_interval = poll_interval
+ self.end_time = end_time
+ self.execution_deadline = execution_deadline
+ self.database = database
+ self.schema = schema
+ self.role = role
+ self.warehouse = warehouse
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize SnowparkContainerJobTrigger arguments and class path."""
+ return (
+
"airflow.providers.snowflake.triggers.snowpark_containers.SnowparkContainerJobTrigger",
+ {
+ "job_name": self.job_name,
+ "snowflake_conn_id": self.snowflake_conn_id,
+ "poll_interval": self.poll_interval,
+ "end_time": self.end_time,
+ "execution_deadline": self.execution_deadline,
+ "database": self.database,
+ "schema": self.schema,
+ "role": self.role,
+ "warehouse": self.warehouse,
+ },
+ )
+
+ def _get_hook(self) -> SnowflakeHook:
+ """Build a ``SnowflakeHook`` from the trigger's connection settings."""
+ return SnowflakeHook(
+ snowflake_conn_id=self.snowflake_conn_id,
+ warehouse=self.warehouse,
+ database=self.database,
+ schema=self.schema,
+ role=self.role,
+ )
+
+ async def _describe_status(self, hook: SnowflakeHook) -> str | None:
+ """Return the job's current status via ``DESCRIBE SERVICE``, or
``None`` if absent."""
+ # SnowflakeHook is synchronous. Run the blocking poll off the event
loop so a
+ # single query does not stall every other trigger on this triggerer.
+ response: Any = await asyncio.to_thread(
+ hook.run,
+ f"DESCRIBE SERVICE {self.job_name}",
+ handler=fetch_one_handler,
+ return_dictionaries=True,
+ )
+ return response.get("status") if response else None
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """Poll the job status and yield exactly one terminal event."""
+ hook = self._get_hook()
+ while True:
+ now = time.time()
+ if self.execution_deadline is not None and now >=
self.execution_deadline:
+ yield TriggerEvent(
+ {
+ "status": "timeout",
+ "job_name": self.job_name,
+ "message": f"Job {self.job_name} reached the execution
timeout.",
+ }
+ )
+ return
+
+ if now >= self.end_time:
+ yield TriggerEvent(
+ {
+ "status": "timeout",
+ "job_name": self.job_name,
+ "message": f"Job {self.job_name} did not reach a
terminal status before the timeout.",
+ }
+ )
+ return
+
+ try:
+ status = await self._describe_status(hook=hook)
+ except Exception as e:
+ yield TriggerEvent({"status": "error", "job_name":
self.job_name, "message": str(e)})
+ return
+
+ if status in TERMINAL_STATUSES:
+ yield TriggerEvent({"status": status, "job_name":
self.job_name})
+ return
+
+ if status not in NON_TERMINAL_STATUSES:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "job_name": self.job_name,
+ "message": f"Job {self.job_name} returned unexpected
status: {status}",
+ }
+ )
+ return
+
+ await asyncio.sleep(self.poll_interval)
+
+ async def on_kill(self) -> None:
Review Comment:
`BaseTrigger.on_kill()` only landed in Airflow 3.3.0 (apache/airflow#65590),
and this provider's floor is `apache-airflow>=2.11.0`. On 2.11 through 3.2
nothing calls this method, so clearing or marking-failed a deferred SPCS task
leaves the job service running and billing — and the operator's own `on_kill`
doesn't cover that either, since a deferred task isn't on a worker.
The code itself is right and forward-compatible; it's the version gap that's
worth surfacing. A sentence in the `deferrable` docstring or in
`docs/operators/snowpark_containers.rst` saying that kill-time cleanup of a
deferred job requires Airflow 3.3+ would save someone on 3.1 an unexplained
orphaned service.
--
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]
