qaziashikin opened a new pull request, #67915:
URL: https://github.com/apache/airflow/pull/67915
Update Notebook Operator to pull S3 project bucket from tooling env
The `SageMakerUnifiedStudioNotebook` operator and sensor read notebook
output variables from the project's S3 bucket. Previously the bucket name was
constructed from a hardcoded template:
f"amazon-sagemaker-{account_id}-{region}-{project_id}"
That template only matches projects whose bucket was provisioned with the
default naming convention. Projects that bring their own bucket (BYOB), use a
custom S3 prefix, or are provisioned through the ToolingLite blueprint have a
different bucket name and the output read silently failed.
Resolve the bucket the same way the SMUS Workflows UI does:
- List the project's "Tooling" environment via DataZone
ListEnvironmentBlueprints + ListEnvironments.
- Fall back to "ToolingLite" when no Tooling environment exists.
- GetEnvironment to read provisionedResources, and parse the bucket
name from the s3BucketPath resource.
The operator and sensor now pass `domain_identifier` through to
`get_notebook_outputs` so the lookup has both identifiers it needs.
Verified end-to-end against a real SMUS domain: notebook outputs are still
read and pushed to xcom for
downstream consumption.
```
IAM:
2026-06-02T16:39:46.209602Z [info ] Reading notebook outputs from
s3://amazon-sagemaker-377228489309-us-west-2-bisglciuqpgv0w/sys/notebooks/bvnn0s06pp42a8/runs/avf69jluff7ki8/notebook_outputs.json
[airflow.providers.amazon.aws.hooks.sagemaker_unified_studio_notebook]
2026-06-02T16:39:47.493299Z [info ] Successfully read 5 notebook
output(s).
[airflow.providers.amazon.aws.hooks.sagemaker_unified_studio_notebook]
IDC:
2026-06-02T17:37:10.992022Z [info ] Exiting notebook run 5jackgb0j3uh0n.
Status: SUCCEEDED
[airflow.task.hooks.airflow.providers.amazon.aws.hooks.sagemaker_unified_studio_notebook.SageMakerUnifiedStudioNotebookHook]
2026-06-02T17:37:12.003754Z [info ] Reading notebook outputs from
s3://amazon-maxdome-430606112922-us-east-1-805540179/sys/notebooks/4h1klahfqowoiv/runs/5jackgb0j3uh0n/notebook_outputs.json
[airflow.providers.amazon.aws.hooks.sagemaker_unified_studio_notebook]
Dag run in success state
Dag run start:2026-06-02 17:33:14.883392+00:00 end:2026-06-02
17:43:12.021628+00:00
2026-06-02T17:43:12.021757Z [info ] DagRun Finished:
dag_id=example_sagemaker_unified_studio_notebook, logical_date=2026-06-02
17:33:14.883392+00:00, run_id=manual__2026-06-02T17:33:14.889842+00:00,
run_start_date=2026-06-02 17:33:14.883392+00:00, run_end_date=2026-06-02
17:43:12.021628+00:00, run_duration=597.138236, state=success, run_type=manual,
data_interval_start=2026-06-02 17:33:14.883392+00:00,
data_interval_end=2026-06-02 17:33:14.883392+00:00,
[airflow.models.dagrun.DagRun]
PASSED
=================== 1 passed, 1 warning in 608.57s (0:10:08)
===================
```
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull
request.
Write a good git commit message following this guide:
http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the
documentation.
Feel free to ping (in general) for the review if you do not see reaction for
a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we
sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR,
please
change below checkbox to `[X]` followed by the name of the tool, uncomment
the "Generated-by".
-->
- [X] Yes (please specify the tool below)
<!--
Generated-by: [Tool Name] following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
--
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]
