Skip to content

Fix SageMakerUnifiedStudioNotebookOperator workflow_name under MWAA Serverless - #72335

Open
qaziashikin wants to merge 1 commit into
apache:mainfrom
qaziashikin:smus-notebook-workflow-name-fix
Open

Fix SageMakerUnifiedStudioNotebookOperator workflow_name under MWAA Serverless#72335
qaziashikin wants to merge 1 commit into
apache:mainfrom
qaziashikin:smus-notebook-workflow-name-fix

Conversation

@qaziashikin

Copy link
Copy Markdown
Contributor

Description

Under MWAA Serverless (Overdrive), dag_id is overwritten with a per-execution UUID for
tenant isolation, so the operator was passing the execution UUID as workflow_name to
StartNotebookRun instead of the actual workflow name. This breaks the SMUS UI's ability
to link a notebook run back to the workflow that triggered it.

MWAA Serverless now injects the workflow name as a DAG param under the key
mwaa_serverless_workflow_id. Read from that param when set, fall back to
context["dag"].dag_id for standard SMUS Workflows and open-source Airflow.

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)
    Claude Opus 4.7

Testing

Ran system test:

qashikin@c889f3b847e4 airflow % .venv/bin/python -m pytest --system -v --with-db-init \
  --setup-timeout=3600 --execution-timeout=3600 --teardown-timeout=3600 \
  providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_notebook.py
======================================== test session starts =========================================
platform darwin -- Python 3.11.14, pytest-9.1.1, pluggy-1.6.0 -- /Volumes/workplace/airflow/.venv/bin/python
cachedir: .pytest_cache
rootdir: /Volumes/workplace/airflow
configfile: pyproject.toml
plugins: mock-3.15.1, instafail-0.5.0, time-machine-3.4.0, unordered-0.8.0, cov-7.1.0, timeouts-1.2.1, xdist-3.8.0, asyncio-1.4.0, custom-exit-code-0.3.0, anyio-4.14.2, icdiff-0.9, rerunfailures-16.6, kgb-7.3, requests-mock-1.12.1
asyncio: mode=Mode.STRICT, debug=False, asyncio_default_fixture_loop_scope=function, asyncio_default_test_loop_scope=function
setup timeout: 3600.0s, execution timeout: 3600.0s, teardown timeout: 3600.0s
collected 1 item                                                                                     

providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_notebook.py::test_runel-common/src/tests_common/test_utils/system_tests.py PASSED [100%]

========================== 1 passed, 1 warning in 290.96s (0:04:50) ===========================

  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@qaziashikin
qaziashikin requested a review from o-nikolas as a code owner August 31, 2026 15:12
@boring-cyborg boring-cyborg Bot added area:providers provider:amazon AWS/Amazon - related issues labels Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:providers provider:amazon AWS/Amazon - related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants