potiuk commented on code in PR #70103:
URL: https://github.com/apache/airflow/pull/70103#discussion_r3890092549
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -203,12 +191,16 @@ def _submit_job(self) -> str:
def _poll_for_status(self) -> str:
"""Poll until the job reaches a terminal state."""
+ end_time = time.time() + self.timeout
while True:
+ if time.time() >= end_time:
+ self._drop_service()
Review Comment:
On timeout the service is dropped before any container logs are fetched, and
a job that hung is exactly the case where those logs are wanted — dropping the
service destroys them. Consider `self._log_container_output(status)` before the
drop here and in the `"timeout"` branch of `execute_complete`. Separately: this
path drops unconditionally, ignoring `drop_on_completion=False`. It's
documented on the `timeout` param so it may well be deliberate, but it's worth
a second look, since a *failed* job is left in place while a *timed-out* one is
not.
##########
providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py:
##########
@@ -64,6 +65,39 @@
from airflow.providers.openlineage.sqlparser import DatabaseInfo
+class SnowparkContainerJobStatus(str, Enum):
Review Comment:
These Snowpark Container Services statuses now live in the SQL hook module,
which is otherwise `SnowflakeHook` plus its OpenLineage surface. I take it the
move is to break the operator↔trigger import cycle —
`triggers/snowpark_containers.py` would do that equally well, since the
operator already imports the trigger and the trigger imports nothing from the
operator. Non-blocking, just a placement thought.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -242,10 +246,47 @@ def execute(self, context: Context) -> str:
raise RuntimeError("Job name was not returned")
if not self.wait_for_completion:
return self.job_name
+ if self.deferrable:
+ # timeout and execution_timeout give the trigger two separate
deadlines. timeout caps
+ # how long the job is polled, and execution_timeout, when set,
enforces the task-level
+ # limit. The trigger times out on whichever is reached first.
+ now = time.time()
+ poll_buffer = timedelta(seconds=self.poll_interval + 60)
+ execution_deadline = None
+ defer_timeout = timedelta(seconds=self.timeout) + poll_buffer
+ if self.execution_timeout is not None:
+ # Hand the execution deadline to the trigger so it emits a
timeout event that drops the
+ # service. The framework's defer timeout would otherwise kill
the task with no cleanup.
+ execution_deadline = now +
self.execution_timeout.total_seconds()
Review Comment:
`now` is `time.time()` at the moment of deferral, so everything before this
point — connection setup, `_submit_job`, the `DESCRIBE` round trip — falls
outside the window. `execution_timeout` is a task-level limit measured from
task start, so a task with `execution_timeout=120` that spent 30s submitting
gets 150s of wall clock. The comment two lines up says this "enforces the
task-level limit", which is what makes it worth fixing rather than documenting.
Deriving the deadline from the task's start instead would be exact:
```suggestion
execution_deadline = (
context["ti"].start_date.timestamp() +
self.execution_timeout.total_seconds()
)
```
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -19,51 +19,25 @@
import time
from collections.abc import Sequence
-from enum import Enum
+from datetime import timedelta
from functools import cached_property
from typing import TYPE_CHECKING, Any
+from airflow.providers.common.compat.sdk import conf
from airflow.providers.common.compat.standard.operators import BaseOperator
from airflow.providers.common.sql.hooks.handlers import fetch_one_handler
-from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
+from airflow.providers.snowflake.hooks.snowflake import (
+ CONTAINER_JOB_NON_TERMINAL_STATUSES,
Review Comment:
`TERMINAL_STATUSES` and `NON_TERMINAL_STATUSES` were public module-level
names in this module and shipped in 6.15.0 (`Add SnowparkContainerJobOperator
(#68259)`), still present in 6.16.1. Renaming them to `CONTAINER_JOB_*` and
moving them to the hook turns `from
airflow.providers.snowflake.operators.snowpark_containers import
TERMINAL_STATUSES` into an `ImportError` on what is otherwise a feature
release. `SnowparkContainerJobStatus` is unaffected — it's re-imported here —
so it's just these two. Two aliases after the import block would keep the old
spelling working.
--
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]
