potiuk commented on code in PR #76106:
URL: https://github.com/apache/airflow/pull/76106#discussion_r4431313717
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -191,14 +211,41 @@ def _run_one(self, sql: str, return_dictionaries: bool =
False) -> Any:
"""Run a single statement that returns one row via
fetch_one_handler."""
return self._hook.run(sql, handler=fetch_one_handler,
return_dictionaries=return_dictionaries)
- def _submit_job(self) -> str:
+ def _describe_status(self, external_id: JsonValue) -> str:
+ """Describe the job's current status."""
+ response = self._run_one(f"DESCRIBE SERVICE {external_id}",
return_dictionaries=True)
+ return response.get("status")
+
+ def submit_job(self, context: Context) -> str:
"""Submit the job and return the name."""
response = self._run_one(self._build_sql())
- job_name = response[0].split("'")[1]
- return job_name
+ self.job_name = response[0].split("'")[1]
+ if not self.job_name:
+ raise RuntimeError("Job name was not returned")
+ return self.job_name
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ """Return NOT_FOUND when the service no longer exists, otherwise the
current status."""
+ try:
+ return self._describe_status(external_id)
+ except ProgrammingError as e:
+ if e.errno == OBJECT_NOT_EXIST_ERROR_CODE:
+ return NOT_FOUND_STATUS
+ raise
+
+ def is_job_active(self, status: str) -> bool:
Review Comment:
`is_job_active` reuses `NON_TERMINAL_STATUSES`, which includes `CANCELLING`
/ `SUSPENDING` / `DELETING`. After a poll timeout the operator drops the
service and raises. If the retry lands while the service is still `DELETING`,
it reconnects, and once the service disappears `_describe_status` raises
`ProgrammingError` 2003. The poll loop doesn't catch that (only
`get_job_status` does), so the retry fails instead of resubmitting. Could
`is_job_active` only treat `PENDING`/`RUNNING` as reconnectable, or could the
poll loop tolerate 2003?
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -191,14 +211,41 @@ def _run_one(self, sql: str, return_dictionaries: bool =
False) -> Any:
"""Run a single statement that returns one row via
fetch_one_handler."""
return self._hook.run(sql, handler=fetch_one_handler,
return_dictionaries=return_dictionaries)
- def _submit_job(self) -> str:
+ def _describe_status(self, external_id: JsonValue) -> str:
+ """Describe the job's current status."""
+ response = self._run_one(f"DESCRIBE SERVICE {external_id}",
return_dictionaries=True)
+ return response.get("status")
+
+ def submit_job(self, context: Context) -> str:
"""Submit the job and return the name."""
response = self._run_one(self._build_sql())
- job_name = response[0].split("'")[1]
- return job_name
+ self.job_name = response[0].split("'")[1]
+ if not self.job_name:
+ raise RuntimeError("Job name was not returned")
+ return self.job_name
+
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ """Return NOT_FOUND when the service no longer exists, otherwise the
current status."""
+ try:
+ return self._describe_status(external_id)
+ except ProgrammingError as e:
+ if e.errno == OBJECT_NOT_EXIST_ERROR_CODE:
+ return NOT_FOUND_STATUS
+ raise
+
+ def is_job_active(self, status: str) -> bool:
+ """Return True while the job is still running."""
+ return status in NON_TERMINAL_STATUSES
+
+ def is_job_succeeded(self, status: str) -> bool:
+ """Return True when the job has completed successfully."""
+ return status == SnowparkContainerJobStatus.DONE
+
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
+ """Poll the job until it reaches a terminal state and handle the final
status."""
+ # On reconnect the mixin skips submit_job, so set the job name from
the external id here.
+ self.job_name = cast("str", external_id)
- def _poll_for_status(self) -> str:
- """Poll until the job reaches a terminal state."""
status = None
end_time = time.monotonic() + self.timeout
Review Comment:
On reconnect `end_time` is recomputed from scratch, so with durable mode
`timeout` limits each try rather than the job as a whole. The job can be polled
for up to `timeout × tries`. Please either document that or persist the
deadline alongside the job name.
##########
providers/snowflake/src/airflow/providers/snowflake/operators/snowpark_containers.py:
##########
@@ -207,14 +254,25 @@ def _poll_for_status(self) -> str:
if self.drop_on_completion:
self._drop_service()
raise TimeoutError(f"Job {self.job_name} did not reach a
terminal status before the timeout.")
- response = self._run_one(f"DESCRIBE SERVICE {self.job_name}",
return_dictionaries=True)
- status = response.get("status")
+ status = self._describe_status(self.job_name)
if status in TERMINAL_STATUSES:
- return status
+ # get_job_result is skipped when the mixin reconnects to a
still-running job, so
+ # finalize here.
+ self._handle_final_status(status=status)
Review Comment:
With the default `drop_on_completion=True`, the service is dropped as soon
as `DONE` is seen, and nothing records that `DONE` first. If the worker dies
after the drop but before the task's success is recorded, the retry sees
`NOT_FOUND` and resubmits, which is the duplicate run durable mode is meant to
prevent. The mixin's `get_job_status` docstring suggests caching the terminal
status in `task_state_store` when the remote resource is ephemeral. Writing a
"done" marker before `_drop_service()`, and checking it in `get_job_status`,
would close this window.
*Flagged by an adversarial reviewer (codex); cross-checked.*
--
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]
