apache-airflow-providers-amazon 9.4.0rc1__py3-none-any.whl → 9.5.0rc1__py3-none-any.whl

airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py

@@ -0,0 +1,188 @@
+# 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.
+
+"""This module contains the Amazon SageMaker Unified Studio Notebook hook."""
+
+from __future__ import annotations
+
+import time
+
+from sagemaker_studio import ClientConfig
+from sagemaker_studio.sagemaker_studio_api import SageMakerStudioAPI
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+from airflow.providers.amazon.aws.utils.sagemaker_unified_studio import is_local_runner
+
+
+class SageMakerNotebookHook(BaseHook):
+    """
+    Interact with Sagemaker Unified Studio Workflows.
+
+    This hook provides a wrapper around the Sagemaker Workflows Notebook Execution API.
+
+    Examples:
+     .. code-block:: python
+
+        from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import SageMakerNotebookHook
+
+        notebook_hook = SageMakerNotebookHook(
+            input_config={"input_path": "path/to/notebook.ipynb", "input_params": {"param1": "value1"}},
+            output_config={"output_uri": "folder/output/location/prefix", "output_formats": "NOTEBOOK"},
+            execution_name="notebook_execution",
+            waiter_delay=10,
+            waiter_max_attempts=1440,
+        )
+
+    :param execution_name: The name of the notebook job to be executed, this is same as task_id.
+    :param input_config: Configuration for the input file.
+        Example: {'input_path': 'folder/input/notebook.ipynb', 'input_params': {'param1': 'value1'}}
+    :param output_config: Configuration for the output format. It should include an output_formats parameter to specify the output format.
+        Example: {'output_formats': ['NOTEBOOK']}
+    :param compute: compute configuration to use for the notebook execution. This is a required attribute
+        if the execution is on a remote compute.
+        Example: { "instance_type": "ml.m5.large", "volume_size_in_gb": 30, "volume_kms_key_id": "", "image_uri": "string", "container_entrypoint": [ "string" ]}
+    :param termination_condition: conditions to match to terminate the remote execution.
+        Example: { "MaxRuntimeInSeconds": 3600 }
+    :param tags: tags to be associated with the remote execution runs.
+        Example: { "md_analytics": "logs" }
+    :param waiter_delay: Interval in seconds to check the task execution status.
+    :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
+    """
+
+    def __init__(
+        self,
+        execution_name: str,
+        input_config: dict | None = None,
+        output_config: dict | None = None,
+        compute: dict | None = None,
+        termination_condition: dict | None = None,
+        tags: dict | None = None,
+        waiter_delay: int = 10,
+        waiter_max_attempts: int = 1440,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self._sagemaker_studio = SageMakerStudioAPI(self._get_sagemaker_studio_config())
+        self.execution_name = execution_name
+        self.input_config = input_config or {}
+        self.output_config = output_config or {"output_formats": ["NOTEBOOK"]}
+        self.compute = compute
+        self.termination_condition = termination_condition or {}
+        self.tags = tags or {}
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+
+    def _get_sagemaker_studio_config(self):
+        config = ClientConfig()
+        config.overrides["execution"] = {"local": is_local_runner()}
+        return config
+
+    def _format_start_execution_input_config(self):
+        config = {
+            "notebook_config": {
+                "input_path": self.input_config.get("input_path"),
+                "input_parameters": self.input_config.get("input_params"),
+            },
+        }
+
+        return config
+
+    def _format_start_execution_output_config(self):
+        output_formats = self.output_config.get("output_formats")
+        config = {
+            "notebook_config": {
+                "output_formats": output_formats,
+            }
+        }
+        return config
+
+    def start_notebook_execution(self):
+        start_execution_params = {
+            "execution_name": self.execution_name,
+            "execution_type": "NOTEBOOK",
+            "input_config": self._format_start_execution_input_config(),
+            "output_config": self._format_start_execution_output_config(),
+            "termination_condition": self.termination_condition,
+            "tags": self.tags,
+        }
+        if self.compute:
+            start_execution_params["compute"] = self.compute
+        else:
+            start_execution_params["compute"] = {"instance_type": "ml.m4.xlarge"}
+
+        print(start_execution_params)
+        return self._sagemaker_studio.execution_client.start_execution(**start_execution_params)
+
+    def wait_for_execution_completion(self, execution_id, context):
+        wait_attempts = 0
+        while wait_attempts < self.waiter_max_attempts:
+            wait_attempts += 1
+            time.sleep(self.waiter_delay)
+            response = self._sagemaker_studio.execution_client.get_execution(execution_id=execution_id)
+            error_message = response.get("error_details", {}).get("error_message")
+            status = response["status"]
+            if "files" in response:
+                self._set_xcom_files(response["files"], context)
+            if "s3_path" in response:
+                self._set_xcom_s3_path(response["s3_path"], context)
+
+            ret = self._handle_state(execution_id, status, error_message)
+            if ret:
+                return ret
+
+        # If timeout, handle state FAILED with timeout message
+        return self._handle_state(execution_id, "FAILED", "Execution timed out")
+
+    def _set_xcom_files(self, files, context):
+        if not context:
+            error_message = "context is required"
+            raise AirflowException(error_message)
+        for file in files:
+            context["ti"].xcom_push(
+                key=f"{file['display_name']}.{file['file_format']}",
+                value=file["file_path"],
+            )
+
+    def _set_xcom_s3_path(self, s3_path, context):
+        if not context:
+            error_message = "context is required"
+            raise AirflowException(error_message)
+        context["ti"].xcom_push(
+            key="s3_path",
+            value=s3_path,
+        )
+
+    def _handle_state(self, execution_id, status, error_message):
+        finished_states = ["COMPLETED"]
+        in_progress_states = ["IN_PROGRESS", "STOPPING"]
+
+        if status in in_progress_states:
+            info_message = f"Execution {execution_id} is still in progress with state:{status}, will check for a terminal status again in {self.waiter_delay}"
+            self.log.info(info_message)
+            return None
+        execution_message = f"Exiting Execution {execution_id} State: {status}"
+        if status in finished_states:
+            self.log.info(execution_message)
+            return {"Status": status, "ExecutionId": execution_id}
+        else:
+            log_error_message = f"Execution {execution_id} failed with error: {error_message}"
+            self.log.error(log_error_message)
+            if error_message == "":
+                error_message = execution_message
+            raise AirflowException(error_message)

airflow/providers/amazon/aws/links/base_aws.py

@@ -19,14 +19,20 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
 
-from airflow.models import BaseOperatorLink, XCom
+from airflow.models import XCom
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
     from airflow.models import BaseOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.utils.context import Context
 
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseOperatorLink
+else:
+    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
+
 
 BASE_AWS_CONSOLE_LINK = "https://console.{aws_domain}"
 

airflow/providers/amazon/aws/links/sagemaker_unified_studio.py

@@ -0,0 +1,27 @@
+# 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
+
+from airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class SageMakerUnifiedStudioLink(BaseAwsLink):
+    """Helper class for constructing Amazon SageMaker Unified Studio Links."""
+
+    name = "Amazon SageMaker Unified Studio"
+    key = "sagemaker_unified_studio"
+    format_str = BASE_AWS_CONSOLE_LINK + "/datazone/home?region={region_name}"

airflow/providers/amazon/aws/log/s3_task_handler.py

@@ -26,6 +26,7 @@ from typing import TYPE_CHECKING
 
 from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
@@ -58,7 +59,8 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
     def hook(self):
         """Returns S3Hook."""
         return S3Hook(
-            aws_conn_id=conf.get("logging", "REMOTE_LOG_CONN_ID"), transfer_config_args={"use_threads": False}
+            aws_conn_id=conf.get("logging", "REMOTE_LOG_CONN_ID"),
+            transfer_config_args={"use_threads": False},
         )
 
     def set_context(self, ti: TaskInstance, *, identifier: str | None = None) -> None:
@@ -73,7 +75,7 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
         is_trigger_log_context = getattr(ti, "is_trigger_log_context", False)
         self.upload_on_close = is_trigger_log_context or not getattr(ti, "raw", None)
         # Clear the file first so that duplicate data is not uploaded
-        # when re-using the same path (e.g. with rescheduled sensors)
+        # when reusing the same path (e.g. with rescheduled sensors)
         if self.upload_on_close:
             with open(self.handler.baseFilename, "w"):
                 pass
@@ -116,12 +118,16 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
         keys = self.hook.list_keys(bucket_name=bucket, prefix=prefix)
         if keys:
             keys = sorted(f"s3://{bucket}/{key}" for key in keys)
-            messages.append("Found logs in s3:")
-            messages.extend(f"  * {key}" for key in keys)
+            if AIRFLOW_V_3_0_PLUS:
+                messages = keys
+            else:
+                messages.append("Found logs in s3:")
+                messages.extend(f"  * {key}" for key in keys)
             for key in keys:
                 logs.append(self.s3_read(key, return_error=True))
         else:
-            messages.append(f"No logs found on s3 for ti={ti}")
+            if not AIRFLOW_V_3_0_PLUS:
+                messages.append(f"No logs found on s3 for ti={ti}")
         return messages, logs
 
     def s3_log_exists(self, remote_log_location: str) -> bool:
@@ -152,7 +158,13 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
                 return msg
         return ""
 
-    def s3_write(self, log: str, remote_log_location: str, append: bool = True, max_retry: int = 1) -> bool:
+    def s3_write(
+        self,
+        log: str,
+        remote_log_location: str,
+        append: bool = True,
+        max_retry: int = 1,
+    ) -> bool:
         """
         Write the log to the remote_log_location; return `True` or fails silently and return `False`.
 
@@ -185,7 +197,10 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
                 break
             except Exception:
                 if try_num < max_retry:
-                    self.log.warning("Failed attempt to write logs to %s, will retry", remote_log_location)
+                    self.log.warning(
+                        "Failed attempt to write logs to %s, will retry",
+                        remote_log_location,
+                    )
                 else:
                     self.log.exception("Could not write logs to %s", remote_log_location)
                     return False