Add AWS SageMaker Unified Studio Workflow Operator #45726

agupta01 · 2025-01-16T21:21:39Z

Description

Adds an operator used for executing Jupyter Notebooks, Querybooks, and Visual ETL jobs within the context of a SageMaker Unified Studio project.

SageMaker Unified Studio (SUS) supports development of Airflow DAGs (called "workflows" within the product) that are run on an MWAA cluster managed by the project. These workflows have the ability to orchestrate the execution of Unified Studio artifacts that can connect to data assets stored in a SUS project.

Implementation-wise, these notebooks are executed on a SageMaker Training Job running a SageMaker Distribution environment within the context of a SUS project.

Components

SageMakerNotebookOperator: this operator allows users to execute Unified Studio artifacts within the context of their project.
SageMakerNotebookHook: this hook provides a wrapper around the notebook execution
SageMakerNotebookSensor: this sensor waits on status updates from the notebook execution
SageMakerNotebookJobTrigger: this trigger activates when the notebook execution completes

Usage

Note that this operator introduces a dependency on the SageMaker Studio SDK https://www.pypi.org/project/sagemaker-studio

with DAG(...) as dag:
    ...
    run_notebook = SageMakerNotebookOperator(
        task_id="initial",
        input_config={"input_path": <notebook_path_in_s3>, "input_params": {}},
        output_config={"output_formats": ["NOTEBOOK"]},
        wait_for_completion=True,
        poll_interval=5,
    )
   ...

Testing

MWAA uses python 3.11 and postgres backend, so we will set those values for all tests.

Unit tests

breeze testing core-tests -p 3.11 -b postgres providers/tests/amazon/aws/*/test_sagemaker_unified_studio.py

System tests

Ensure a properly configured SageMaker Unified Domain and Project as indicated in the example_sagemaker_unified_studio.py file. Also ensure AWS credentials are populated and up to date. Then, populate the DOMAIN_ID, PROJECT_ID, ENVIRONMENT_ID, and S3_PATH in files/airflow-breeze-config/variables.env and run:

breeze testing system-tests -p 3.11 -b postgres --forward-credentials --test-timeout 600 providers/tests/system/amazon/aws/example_sagemaker_unified_studio.py

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

… hooks

boring-cyborg · 2025-01-16T21:21:44Z

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:

Pay attention to the quality of your code (ruff, mypy and type annotations). Our pre-commits will help you with that.
In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example DAG that shows how users should use it.
Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
Be sure to read the Airflow Coding style.
Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: [email protected]
Slack: https://s.apache.org/airflow-slack

o-nikolas · 2025-01-16T21:38:36Z

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

+    :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 control


This sentence seems to just tail off in the middle?

o-nikolas · 2025-01-16T21:39:03Z

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

+        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 control
+        Example: {'output_formats': ['NOTEBOOK']}
+    :param compute: compute configuration to use for the notebook execution. This is an required attribute


Suggested change

:param compute: compute configuration to use for the notebook execution. This is an required attribute

:param compute: compute configuration to use for the notebook execution. This is a required attribute

o-nikolas · 2025-01-16T21:44:18Z

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

+
+    def _format_start_execution_output_config(self):
+        output_formats = (
+            self.output_config.get("output_formats") if self.output_config else ["NOTEBOOK"]


This ternary is unnecessary right? There is a default value provided in the constructor so the output_config can't be empty?

o-nikolas · 2025-01-16T21:50:43Z

providers/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py

+    :param task_id: A unique, meaningful id for the task.
+    :param input_config: Configuration for the input file. Input path should be specified as a relative path.
+        The provided relative path will be automatically resolved to an absolute path within
+        the context of the user's home directory in the IDE. Input parms should be a dict


Suggested change

the context of the user's home directory in the IDE. Input parms should be a dict

the context of the user's home directory in the IDE. Input parsms should be a dict

o-nikolas · 2025-01-28T23:15:01Z

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

+        return self._sagemaker_studio.execution_client.start_execution(request)
+
+    def wait_for_execution_completion(self, execution_id, context):
+


Suggested change

o-nikolas · 2025-01-28T23:27:46Z

providers/src/airflow/providers/amazon/aws/sensors/sagemaker_unified_studio.py

+from airflow.utils.context import Context
+
+
+class SageMakerNotebookSensor(BaseSensorOperator):


Use AwsBaseSensor?

o-nikolas · 2025-01-28T23:34:59Z

providers/tests/system/amazon/aws/example_sagemaker_unified_studio.py

+    Username: airflowTestUser
+    Password: airflowSystemTestP@ssword1!


I don't see these credentials used anywhere in the test? Is it really mandatory? If yes, we should make it configurable and provide them through the test context builder

o-nikolas · 2025-01-28T23:55:58Z

providers/tests/system/amazon/aws/example_sagemaker_unified_studio.py

+@task
+def emulate_mwaa_environment(
+    domain_id: str, project_id: str, environment_id: str, s3_path: str
+):


If run using an container based executor (like ECS or K8s) this will have no effect. A container will spin up, export these envs, then it will get torn down and the next task will run in a new container. So this test will fail to run on our ECS executor test suite.

Any other way around this? Otherwise we'll need to create an image for container based tests, or at least provide these env vars through executor_config

o-nikolas · 2025-01-28T23:56:23Z

providers/tests/system/amazon/aws/example_sagemaker_unified_studio.py

+    notebook_path = "test_notebook.ipynb"  # This should be the path to your .ipynb, .sqlnb, or .vetl file in your project.
+
+    run_notebook = SageMakerNotebookOperator(
+        task_id="initial",


Somewhat strange task_id

o-nikolas · 2025-01-28T23:57:09Z

providers/tests/system/amazon/aws/example_sagemaker_unified_studio.py

+        run_notebook,
+    )


No teardown needed? What if things break, get stuck, timeout? Any way to manually stop a noteboook execution?

agupta01 and others added 9 commits January 13, 2025 20:35

Add sagemaker_unified_studio notebook operator, sensor, triggers, and…

1118295

… hooks

Fix sagemaker_unified_studio unit tests

44e45c8

Add basic system test for sagemaker_unified_studio

6c627ef

Add setup/teardown stubs for sagemaker_unified_studio system test

33c4a3d

Add more specifics to SUS system test

a1e0766

Update name of SUS helper

5ee6aad

Cleanup and format SUS system test

e378358

Merge branch 'apache:main' into main

b369756

Merge branch 'apache:main' into main

209cb3c

boring-cyborg bot added area:providers provider:amazon AWS/Amazon - related issues labels Jan 16, 2025

agupta01 and others added 8 commits January 17, 2025 12:10

Merge branch 'apache:main' into main

092c786

Merge branch 'apache:main' into main

55f5486

Fix notebook path in SUS system test

c640efa

Merge branch 'apache:main' into main

41bd997

Merge branch 'apache:main' into main

7acc649

Update SUS docs

b7ad9db

Update SUS docs to include vETL

7ca5022

Clarity updates on SUS docs

13a10ce

o-nikolas reviewed Jan 28, 2025

View reviewed changes

agupta01 changed the title ~~Add AWS SageMaker Unified Studio Notebook Operator~~ Add AWS SageMaker Unified Studio Workflow Operator Jan 29, 2025

Merge branch 'apache:main' into main

756fcad

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add AWS SageMaker Unified Studio Workflow Operator #45726

Add AWS SageMaker Unified Studio Workflow Operator #45726

agupta01 commented Jan 16, 2025 •

edited

Loading

boring-cyborg bot commented Jan 16, 2025

o-nikolas Jan 16, 2025

o-nikolas Jan 16, 2025

o-nikolas Jan 16, 2025

o-nikolas Jan 16, 2025

o-nikolas Jan 28, 2025

o-nikolas Jan 28, 2025

o-nikolas Jan 28, 2025

o-nikolas Jan 28, 2025

o-nikolas Jan 28, 2025

o-nikolas Jan 28, 2025

	:param compute: compute configuration to use for the notebook execution. This is an required attribute
	:param compute: compute configuration to use for the notebook execution. This is a required attribute

	the context of the user's home directory in the IDE. Input parms should be a dict
	the context of the user's home directory in the IDE. Input parsms should be a dict

		return self._sagemaker_studio.execution_client.start_execution(request)

		def wait_for_execution_completion(self, execution_id, context):

		from airflow.utils.context import Context


		class SageMakerNotebookSensor(BaseSensorOperator):

		Username: airflowTestUser
		Password: airflowSystemTestP@ssword1!

Add AWS SageMaker Unified Studio Workflow Operator #45726

Are you sure you want to change the base?

Add AWS SageMaker Unified Studio Workflow Operator #45726

Conversation

agupta01 commented Jan 16, 2025 • edited Loading

Description

Components

Usage

Testing

Unit tests

System tests

boring-cyborg bot commented Jan 16, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

agupta01 commented Jan 16, 2025 •

edited

Loading