o-nikolas commented on code in PR #65571: URL: https://github.com/apache/airflow/pull/65571#discussion_r3125699625
########## providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py: ########## @@ -7,60 +7,100 @@ # 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, Review Comment: Why did you remove part of the license? This will fail pre checks ########## providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py: ########## @@ -7,60 +7,100 @@ # 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. - -"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger.""" from __future__ import annotations -from airflow.triggers.base import BaseTrigger +import asyncio +from typing import Any, Dict Review Comment: You don't need to import Dict for typing, just use the regular builtin ########## providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker_unified_studio.py: ########## @@ -7,60 +7,100 @@ # 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. - -"""This module contains the Amazon SageMaker Unified Studio Notebook job trigger.""" from __future__ import annotations -from airflow.triggers.base import BaseTrigger +import asyncio +from typing import Any, Dict + +from airflow.triggers.base import BaseTrigger, TriggerEvent +from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import ( + SageMakerNotebookHook, +) class SageMakerNotebookJobTrigger(BaseTrigger): """ - Watches for a notebook job, triggers when it finishes. - - Examples: - .. code-block:: python - - from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import SageMakerNotebookJobTrigger + Async trigger that polls a SageMaker Unified Studio notebook execution until it + reaches a terminal state. - notebook_trigger = SageMakerNotebookJobTrigger( - execution_id="notebook_job_1234", - execution_name="notebook_task", - waiter_delay=10, - waiter_max_attempts=1440, - ) - - :param execution_id: A unique, meaningful id for the task. - :param execution_name: A unique, meaningful name for the task. - :param waiter_delay: Interval in seconds to check the notebook execution status. - :param waiter_max_attempts: Number of attempts to wait before returning FAILED. + :param execution_id: The execution ID returned by start_notebook_execution(). + :param execution_name: The name of the execution (task_id). + :param waiter_delay: Seconds between polling attempts. + :param waiter_max_attempts: Maximum polling attempts before timing out. """ - def __init__(self, execution_id, execution_name, waiter_delay, waiter_max_attempts, **kwargs): + def __init__( + self, + *, + execution_id: str, + execution_name: str, + waiter_delay: int, + waiter_max_attempts: int, + **kwargs, + ): super().__init__(**kwargs) self.execution_id = execution_id self.execution_name = execution_name self.waiter_delay = waiter_delay self.waiter_max_attempts = waiter_max_attempts - def serialize(self): + def serialize(self) -> tuple[str, Dict[str, Any]]: return ( - # dynamically generate the fully qualified name of the class - self.__class__.__module__ + "." + self.__class__.__qualname__, + "airflow.providers.amazon.aws.triggers.sagemaker_unified_studio.SageMakerNotebookJobTrigger", { "execution_id": self.execution_id, "execution_name": self.execution_name, - "poll_interval": self.poll_interval, + "waiter_delay": self.waiter_delay, + "waiter_max_attempts": self.waiter_max_attempts, }, ) async def run(self): - pass + hook = SageMakerNotebookHook(execution_name=self.execution_name) + loop = asyncio.get_running_loop() + + attempts = 0 + + while attempts < self.waiter_max_attempts: + attempts += 1 + + response = await loop.run_in_executor( + None, + hook._sagemaker_studio.execution_client.get_execution, Review Comment: This is accessing the "private" client attribute. Maybe add a little helper function to the hook class that encapsulates this. -- 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]
