justinpakzad commented on code in PR #70103: URL: https://github.com/apache/airflow/pull/70103#discussion_r3952908139
########## 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: Good point. Added a note in the `snowpark_containers.rst`. ########## 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: Good call out. The switch to best-effort was intentional and I'm aware it does change the current behavior but I think it's the better approach here. A job that already succeeded could be marked as failed (due to a transient error while dropping) and end up retrying an already successful job. I put a comment explaining why. -- 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]
