ruijiang-rjian commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2931760949
##########
providers/amazon/docs/index.rst:
##########
@@ -123,7 +123,7 @@ PIP package Version required
``asgiref`` ``>=2.3.0``
``PyAthena`` ``>=3.10.0``
``jmespath`` ``>=0.7.0``
-``sagemaker-studio`` ``>=1.0.9``
+``sagemaker-studio`` ``>=1.0.25,<1.1.0``
Review Comment:
looks good! although as discussed we need to take a note on our end that we
will make sure our sagemaker-studio changes also go into 1.1.x so that later we
can maybe relax this `<1.1.0`
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -43,7 +43,10 @@
This test will emulate a DAG run in the shared MWAA environment inside a
SageMaker Unified Studio Project.
The setup tasks will set up the project and configure the test runner to
emulate an MWAA instance.
-Then, the SageMakerNotebookOperator will run a test notebook. This should spin
up a SageMaker training job, run the notebook, and exit successfully.
+Then, the SageMakerNotebookOperator will run a test notebook twice:
+ 1. With all environment variables set (legacy MWAA-style, for SDK backwards
compatibility)
+ 2. With domain_id/project_id passed explicitly as operator parameters and no
environment variables
Review Comment:
nit: no environment variables set in the airflow environment
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -156,12 +161,35 @@ def mock_mwaa_environment(parameters: dict):
)
# [END howto_operator_sagemaker_unified_studio_notebook]
+ # [START howto_operator_sagemaker_unified_studio_notebook_explicit_params]
+ # Run notebook with domain_id/project_id passed explicitly as operator
parameters.
+ # No environment variables needed — the SDK resolves the S3 path and
region from domain_id and project_id.
+ run_notebook_explicit_params = SageMakerNotebookOperator(
+ task_id="run-notebook-explicit",
+ domain_id=domain_id,
+ project_id=project_id,
+ input_config={"input_path": notebook_path, "input_params": {}},
+ output_config={"output_formats": ["NOTEBOOK"]}, # optional
+ compute={
+ "instance_type": "ml.m5.large",
+ "volume_size_in_gb": 30,
+ }, # optional
+ termination_condition={"max_runtime_in_seconds": 600}, # optional
+ tags={}, # optional
+ wait_for_completion=True, # optional
+ waiter_delay=5, # optional
+ deferrable=False, # optional
+ )
+ # [END howto_operator_sagemaker_unified_studio_notebook_explicit_params]
+
chain(
# TEST SETUP
test_context,
setup_mwaa_environment,
- # TEST BODY
+ # TEST BODY: legacy env-var-based resolution
run_notebook,
+ # TEST BODY: explicit params, no env vars required
+ run_notebook_explicit_params,
Review Comment:
question: I wonder if we keep this order, isn't the env variables already
being set before `run_notebook_explicit_params`? so when
`run_notebook_explicit_params` gets run, we cannot guarantee the env variables
are not there? or am I missing sth here..
--
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]
