Skip to content

Fix SageMakerNotebookOperator domain_id/project_id/domain_region params#62962

Merged
o-nikolas merged 8 commits intoapache:mainfrom
Shnekit:fix/sagemaker-notebook-operator
Mar 17, 2026
Merged

Fix SageMakerNotebookOperator domain_id/project_id/domain_region params#62962
o-nikolas merged 8 commits intoapache:mainfrom
Shnekit:fix/sagemaker-notebook-operator

Conversation

@Shnekit
Copy link
Copy Markdown
Contributor

@Shnekit Shnekit commented Mar 5, 2026 •
edited
Loading

Summary

This PR fixes two issues with the SageMakerNotebookOperator and its underlying hook when domain_id, project_id, and domain_region are passed as parameters.

Fixes

1. Method _get_sagemaker_studio_config in hook did not specify domain, project, and DZ region identifiers

The hook was missing domain_identifier, project_identifier, and datazone_domain_region.

2. Add template_fields to SageMakerNotebookOperator

Without template_fields, passing XCom references (e.g. task_instance.xcom_pull(...)) for domain_id, project_id, or domain_region caused a ParamValidationError because the hook was instantiated before Airflow resolved the XCom values to actual strings. Adding these to template_fields ensures they are rendered before execute() is called.

System Tests

  • Updated example_sagemaker_unified_studio to remove deprecated env vars (ENVIRONMENT_ID, SCOPE_NAME, STAGE, ENDPOINT)
  • Added example_sagemaker_unified_studio_explicit_params — identical flow but passes domain_id, project_id, domain_region directly as operator parameters instead of environment variables

@Shnekit Shnekit requested a review from o-nikolas as a code owner March 5, 2026 21:19
@boring-cyborg boring-cyborg bot added area:providers provider:amazon AWS/Amazon - related issues labels Mar 5, 2026
ruijiang-rjian
ruijiang-rjian reviewed Mar 5, 2026
View reviewed changes
@Shnekit
Copy link
Copy Markdown
Contributor Author

Shnekit commented Mar 9, 2026

Tests are still failing due to the fact that sagemaker-studio v1.0.25 is not available yet. Waiting for it to become live to rerun the tests

uranusjr
uranusjr approved these changes Mar 10, 2026
View reviewed changes
Copy link
Copy Markdown
Member

@uranusjr uranusjr left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No idea on how SageMaker works but Python-wise this looks as reasonable as it goes.

@uranusjr
Copy link
Copy Markdown
Member

uranusjr commented Mar 10, 2026

Is the version upgrade needed? It’s not explained above if this uses new features.

@potiuk potiuk marked this pull request as draft March 12, 2026 00:43
@potiuk
Copy link
Copy Markdown
Member

potiuk commented Mar 12, 2026

@Shnekit This PR has been converted to draft because it does not yet meet our Pull Request quality criteria.

Issues found:

  • ❌ Pre-commit / static checks: Failing: CI image checks / Static checks. Run prek run --from-ref main locally to find and fix issues. See Pre-commit / static checks docs.
  • ❌ Provider tests: Failing: Special tests / Min SQLAlchemy test: providers / DB-prov:MinSQLAlchemy-Postgres:14:3.10:-amazon,celer...standard, Special tests / Latest SQLAlchemy test: providers / DB-prov:LatestSQLAlchemy-Postgres:14:3.10:-amazon,celer...standard. Run provider tests with breeze run pytest <provider-test-path> -xvs. See Provider tests docs.
  • ⚠️ Unresolved review comments: This PR has 3 unresolved review threads from maintainers. Please review and resolve all inline review comments before requesting another review. You can resolve a conversation by clicking 'Resolve conversation' on each thread after addressing the feedback. See pull request guidelines.

Note: Your branch is 9 commits behind main. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.

What to do next:

  • The comment informs you what you need to do.
  • Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
  • Maintainers will then proceed with a normal review.

Converting a PR to draft is not a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. If you have questions, feel free to ask on the Airflow Slack.

@Shnekit
Copy link
Copy Markdown
Contributor Author

Shnekit commented Mar 12, 2026

Tightening the sagemaker-studio version to make sure it supports receiving domain id and project id parameters configured through SageMakerNotebookOperator
cc: @uranusjr

@Shnekit Shnekit marked this pull request as ready for review March 13, 2026 14:45
ruijiang-rjian
ruijiang-rjian reviewed Mar 13, 2026
View reviewed changes
@Shnekit Shnekit requested a review from ruijiang-rjian March 13, 2026 17:11
eladkal
eladkal reviewed Mar 16, 2026
View reviewed changes
@Shnekit Shnekit force-pushed the fix/sagemaker-notebook-operator branch from 56f2e6b to 32a4e86 Compare March 16, 2026 22:16
Shnekit added 8 commits March 17, 2026 07:41
@Shnekit
Fix SageMakerNotebookOperator domain_id/project_id/domain_region params
75f855b 
- Fix duplicate _get_sagemaker_studio_config method in hook that was
  missing domain_identifier, project_identifier, datazone_domain_region
- Add template_fields to SageMakerNotebookOperator so XCom references
  are resolved before execute() is called and the hook is instantiated
- Update example_sagemaker_unified_studio system test to remove
  deprecated ENVIRONMENT_ID, SCOPE_NAME, STAGE, ENDPOINT env vars
- Add example_sagemaker_unified_studio_explicit_params system test
  that passes domain_id, project_id, domain_region as operator params
  instead of relying on environment variables
@Shnekit
Fix SageMakerNotebookOperator domain_id/project_id/domain_region params
0749e85 
- Fix duplicate _get_sagemaker_studio_config method in hook that was
  missing domain_identifier, project_identifier, datazone_domain_region
- Add template_fields to SageMakerNotebookOperator so XCom references
  resolve before execute() is called and the hook is instantiated
- Add explicit domain_id/project_id operator params test to system test,
  verifying SDK can resolve S3 path and region without env vars
- Bump sagemaker-studio minimum version to 1.0.25
@Shnekit
Restrict sagemaker-studio version to >=1.0.25,<1.1.0
a47804e
@Shnekit
Fix SageMakerNotebookHook region config and update system test with e…
791c8e1 
…xplicit params path
@Shnekit
Update uv.lock for sagemaker-studio>=1.0.25,<1.1.0 constraint
e4d7830
@Shnekit
Remove pydynamodb dependency and clean up comments
8745a02
@Shnekit
Add explicit params example to SageMaker Unified Studio operator docs
16b3c00
@Shnekit
Fix RST table separator width in docs/index.rst
94a3198
@Shnekit Shnekit force-pushed the fix/sagemaker-notebook-operator branch from 32a4e86 to 94a3198 Compare March 17, 2026 14:43
@Shnekit Shnekit mentioned this pull request Mar 17, 2026
Closed
@o-nikolas o-nikolas merged commit 7f9c310 into apache:main Mar 17, 2026
106 checks passed
imrichardwu pushed a commit to imrichardwu/airflow that referenced this pull request Mar 18, 2026
@Shnekit @imrichardwu
Fix SageMakerNotebookOperator domain_id/project_id/domain_region para…
1188c56 
…ms (apache#62962)

- Method _get_sagemaker_studio_config in hook did not specify domain, project, and DZ region identifiers. The hook was missing domain_identifier, project_identifier, and datazone_domain_region.
- Add template_fields to SageMakerNotebookOperator. Without template_fields, passing XCom references (e.g. task_instance.xcom_pull(...)) for domain_id, project_id, or domain_region caused a ParamValidationError because the hook was instantiated before Airflow resolved the XCom values to actual strings. Adding these to template_fields ensures they are rendered before execute() is called.
- Updated example_sagemaker_unified_studio to reflect changes.
imrichardwu pushed a commit to imrichardwu/airflow that referenced this pull request Mar 18, 2026
@Shnekit @imrichardwu
Fix SageMakerNotebookOperator domain_id/project_id/domain_region para…
af6f42d 
…ms (apache#62962)

- Method _get_sagemaker_studio_config in hook did not specify domain, project, and DZ region identifiers. The hook was missing domain_identifier, project_identifier, and datazone_domain_region.
- Add template_fields to SageMakerNotebookOperator. Without template_fields, passing XCom references (e.g. task_instance.xcom_pull(...)) for domain_id, project_id, or domain_region caused a ParamValidationError because the hook was instantiated before Airflow resolved the XCom values to actual strings. Adding these to template_fields ensures they are rendered before execute() is called.
- Updated example_sagemaker_unified_studio to reflect changes.
fat-catTW pushed a commit to fat-catTW/airflow that referenced this pull request Mar 22, 2026
@Shnekit @fat-catTW
Fix SageMakerNotebookOperator domain_id/project_id/domain_region para…
310194c 
…ms (apache#62962)

- Method _get_sagemaker_studio_config in hook did not specify domain, project, and DZ region identifiers. The hook was missing domain_identifier, project_identifier, and datazone_domain_region.
- Add template_fields to SageMakerNotebookOperator. Without template_fields, passing XCom references (e.g. task_instance.xcom_pull(...)) for domain_id, project_id, or domain_region caused a ParamValidationError because the hook was instantiated before Airflow resolved the XCom values to actual strings. Adding these to template_fields ensures they are rendered before execute() is called.
- Updated example_sagemaker_unified_studio to reflect changes.
techcodie pushed a commit to techcodie/airflow that referenced this pull request Mar 23, 2026
@Shnekit @techcodie
Fix SageMakerNotebookOperator domain_id/project_id/domain_region para…
e3a914b 
…ms (apache#62962)

- Method _get_sagemaker_studio_config in hook did not specify domain, project, and DZ region identifiers. The hook was missing domain_identifier, project_identifier, and datazone_domain_region.
- Add template_fields to SageMakerNotebookOperator. Without template_fields, passing XCom references (e.g. task_instance.xcom_pull(...)) for domain_id, project_id, or domain_region caused a ParamValidationError because the hook was instantiated before Airflow resolved the XCom values to actual strings. Adding these to template_fields ensures they are rendered before execute() is called.
- Updated example_sagemaker_unified_studio to reflect changes.
@jscheffl jscheffl mentioned this pull request Mar 24, 2026
Closed
abhijeets25012-tech pushed a commit to abhijeets25012-tech/airflow that referenced this pull request Apr 9, 2026
@Shnekit
Fix SageMakerNotebookOperator domain_id/project_id/domain_region para…
68bb7b9 
…ms (apache#62962)

- Method _get_sagemaker_studio_config in hook did not specify domain, project, and DZ region identifiers. The hook was missing domain_identifier, project_identifier, and datazone_domain_region.
- Add template_fields to SageMakerNotebookOperator. Without template_fields, passing XCom references (e.g. task_instance.xcom_pull(...)) for domain_id, project_id, or domain_region caused a ParamValidationError because the hook was instantiated before Airflow resolved the XCom values to actual strings. Adding these to template_fields ensures they are rendered before execute() is called.
- Updated example_sagemaker_unified_studio to reflect changes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@eladkal eladkal eladkal left review comments

@uranusjr uranusjr uranusjr approved these changes

@o-nikolas o-nikolas o-nikolas approved these changes

+1 more reviewer

@ruijiang-rjian ruijiang-rjian ruijiang-rjian approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@Shnekit @uranusjr @potiuk @eladkal @o-nikolas @ruijiang-rjian