Add AWS SageMaker Unified Studio Workflow Operator

docs/spelling_wordlist.txt

@@ -319,6 +319,7 @@ connectTimeoutMS
 connexion
 containerConfiguration
 containerd
+ContainerEntrypoint
 ContainerGroup
 containerinstance
 ContainerPort
@@ -693,6 +694,7 @@ Gantt
 gantt
 gapic
 gapped
+gb
 gbq
 gcc
 gcloud
@@ -826,6 +828,7 @@ ImageAnnotatorClient
 imageORfile
 imagePullPolicy
 imagePullSecrets
+ImageUri
 imageVersion
 Imap
 imap
@@ -859,6 +862,7 @@ InstanceFlexibilityPolicy
 InstanceGroupConfig
 InstanceSelection
 instanceTemplates
+InstanceType
 instantiation
 integrations
 interdependencies
@@ -876,6 +880,7 @@ IPv4
 ipv4
 IPv6
 ipv6
+ipynb
 iPython
 irreproducible
 IRSA
@@ -1050,6 +1055,7 @@ masterType
 Matomo
 matomo
 Maxime
+MaxRuntimeInSeconds
 mb
 md
 mediawiki
@@ -1373,6 +1379,8 @@ Qubole
 qubole
 QuboleCheckHook
 Quboles
+querybook
+Querybooks
 queryParameters
 querystring
 queueing
@@ -1887,8 +1895,10 @@ views
 virtualenv
 virtualenvs
 vm
+VolumeKmsKeyId
 VolumeMount
 volumeMounts
+VolumeSizeInGB
 vpc
 WaiterModel
 wape

generated/provider_dependencies.json

@@ -38,6 +38,7 @@
       "jsonpath_ng>=1.5.3",
       "python3-saml>=1.16.0",
       "redshift_connector>=2.0.918",
+      "sagemaker-studio>=1.0.9",
       "watchtower>=3.0.0,!=3.3.0,<4"
     ],
     "devel-deps": [

providers/amazon/README.rst

@@ -67,6 +67,7 @@ PIP package                                 Version required
 ``PyAthena``                                ``>=3.0.10``
 ``jmespath``                                ``>=0.7.0``
 ``python3-saml``                            ``>=1.16.0``
+``sagemaker-studio``                        ``>=1.0.9``
 ==========================================  ======================
 
 Cross provider package dependencies

providers/amazon/docs/operators/sagemakerunifiedstudio.rst

@@ -0,0 +1,60 @@
+ .. 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.
+
+===============================
+Amazon SageMaker Unified Studio
+===============================
+
+`Amazon SageMaker Unified Studio <https://aws.amazon.com/sagemaker/unified-studio/>`__ is a unified development experience that
+brings together AWS data, analytics, artificial intelligence (AI), and machine learning (ML) services.
+It provides a place to build, deploy, execute, and monitor end-to-end workflows from a single interface.
+This helps drive collaboration across teams and facilitate agile development.
+
+Airflow provides operators to orchestrate Notebooks, Querybooks, and Visual ETL jobs within SageMaker Unified Studio Workflows.
+
+Prerequisite Tasks
+------------------
+
+To use these operators, you must do a few things:
+
+  * Create a SageMaker Unified Studio domain and project, following the instruction in `AWS documentation <https://docs.aws.amazon.com/sagemaker-unified-studio/latest/userguide/getting-started.html>`__.
+  * Within your project:
+    * Navigate to the "Compute > Workflow environments" tab, and click "Create" to create a new MWAA environment.
+    * Create a Notebook, Querybook, or Visual ETL job and save it to your project.
+
+Operators
+---------
+
+.. _howto/operator:SageMakerNotebookOperator:
+
+Create an Amazon SageMaker Unified Studio Workflow
+==================================================
+
+To create an Amazon SageMaker Unified Studio workflow to orchestrate your notebook, querybook, and visual ETL runs you can use
+:class:`~airflow.providers.amazon.aws.operators.sagemaker_unified_studio.SageMakerNotebookOperator`.
+
+.. exampleinclude:: /../../providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_sagemaker_unified_studio_notebook]
+    :end-before: [END howto_operator_sagemaker_unified_studio_notebook]
+
+
+Reference
+---------
+
+* `What is Amazon SageMaker Unified Studio <https://docs.aws.amazon.com/sagemaker-unified-studio/latest/userguide/what-is-sagemaker-unified-studio.html>`__

providers/amazon/provider.yaml

@@ -234,6 +234,12 @@ integrations:
     how-to-guide:
       - /docs/apache-airflow-providers-amazon/operators/sagemaker.rst
     tags: [aws]
+  - integration-name: Amazon SageMaker Unified Studio
+    external-doc-url: https://aws.amazon.com/sagemaker/unified-studio/
+    logo: /docs/integration-logos/[email protected]
+    how-to-guide:
+      - /docs/apache-airflow-providers-amazon/operators/sagemakerunifiedstudio.rst
+    tags: [aws]
   - integration-name: Amazon SecretsManager
     external-doc-url: https://aws.amazon.com/secrets-manager/
     logo: /docs/integration-logos/[email protected]
@@ -402,6 +408,9 @@ operators:
   - integration-name: Amazon SageMaker
     python-modules:
       - airflow.providers.amazon.aws.operators.sagemaker
+  - integration-name: Amazon SageMaker Unified Studio
+    python-modules:
+      - airflow.providers.amazon.aws.operators.sagemaker_unified_studio
   - integration-name: Amazon Simple Notification Service (SNS)
     python-modules:
       - airflow.providers.amazon.aws.operators.sns
@@ -503,6 +512,9 @@ sensors:
   - integration-name: Amazon SageMaker
     python-modules:
       - airflow.providers.amazon.aws.sensors.sagemaker
+  - integration-name: Amazon SageMaker Unified Studio
+    python-modules:
+      - airflow.providers.amazon.aws.sensors.sagemaker_unified_studio
   - integration-name: Amazon Simple Queue Service (SQS)
     python-modules:
       - airflow.providers.amazon.aws.sensors.sqs
@@ -627,6 +639,9 @@ hooks:
   - integration-name: Amazon SageMaker
     python-modules:
       - airflow.providers.amazon.aws.hooks.sagemaker
+  - integration-name: Amazon SageMaker Unified Studio
+    python-modules:
+      - airflow.providers.amazon.aws.hooks.sagemaker_unified_studio
   - integration-name: Amazon Simple Email Service (SES)
     python-modules:
       - airflow.providers.amazon.aws.hooks.ses
@@ -699,6 +714,9 @@ triggers:
   - integration-name: Amazon SageMaker
     python-modules:
       - airflow.providers.amazon.aws.triggers.sagemaker
+  - integration-name: Amazon SageMaker Unified Studio
+    python-modules:
+      - airflow.providers.amazon.aws.triggers.sagemaker_unified_studio
   - integration-name: AWS Glue
     python-modules:
       - airflow.providers.amazon.aws.triggers.glue
@@ -734,7 +752,6 @@ triggers:
     python-modules:
       - airflow.providers.amazon.aws.triggers.dms
 
-
 transfers:
   - source-integration-name: Amazon DynamoDB
     target-integration-name: Amazon Simple Storage Service (S3)
@@ -837,6 +854,7 @@ extra-links:
   - airflow.providers.amazon.aws.links.glue.GlueJobRunDetailsLink
   - airflow.providers.amazon.aws.links.logs.CloudWatchEventsLink
   - airflow.providers.amazon.aws.links.sagemaker.SageMakerTransformJobLink
+  - airflow.providers.amazon.aws.links.sagemaker_unified_studio.SageMakerUnifiedStudioLink
   - airflow.providers.amazon.aws.links.step_function.StateMachineDetailsLink
   - airflow.providers.amazon.aws.links.step_function.StateMachineExecutionsDetailsLink
   - airflow.providers.amazon.aws.links.comprehend.ComprehendPiiEntitiesDetectionLink

providers/amazon/pyproject.toml

@@ -75,6 +75,7 @@ dependencies = [
     "PyAthena>=3.0.10",
     "jmespath>=0.7.0",
     "python3-saml>=1.16.0",
+    "sagemaker-studio>=1.0.9",
 ]
 
 # The optional dependencies should be modified in place in the generated file