You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "Shnekit (via GitHub)" <gi...@apache.org> on 2026/03/05 21:19:42 UTC
[PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Shnekit opened a new pull request, #62962:
URL: https://github.com/apache/airflow/pull/62962
## 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. Duplicate `_get_sagemaker_studio_config` in hook**
The hook had a duplicate method — the first (incomplete) version only set `local`, missing `domain_identifier`, `project_identifier`, and `datazone_domain_region`. The duplicate has been removed, leaving only the correct implementation.
**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
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2897911803
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
Review Comment:
Unresolving this one
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2927776821
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
Review Comment:
Fixed
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2927784896
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
+2. A SageMaker Unified Studio Domain (with default VPC and roles)
+3. A project within the SageMaker Unified Studio Domain
+4. A notebook (test_notebook.ipynb) placed in the project's s3 path
+
+This test is identical to example_sagemaker_unified_studio, but passes domain_id, project_id,
+and domain_region explicitly as operator parameters instead of relying on environment variables.
+"""
+
+DAG_ID = "example_sagemaker_unified_studio_explicit_params"
+
+# Externally fetched variables:
+DOMAIN_ID_KEY = "DOMAIN_ID"
+PROJECT_ID_KEY = "PROJECT_ID"
+S3_PATH_KEY = "S3_PATH"
+REGION_NAME_KEY = "REGION_NAME"
+
+sys_test_context_task = (
+ SystemTestContextBuilder()
+ .add_variable(DOMAIN_ID_KEY)
+ .add_variable(PROJECT_ID_KEY)
+ .add_variable(S3_PATH_KEY)
+ .add_variable(REGION_NAME_KEY)
+ .build()
+)
+
+
+def get_mwaa_environment_params(s3_path: str):
+ AIRFLOW_PREFIX = "AIRFLOW__WORKFLOWS__"
+ return {f"{AIRFLOW_PREFIX}PROJECT_S3_PATH": s3_path}
+
+
+@task
+def mock_mwaa_environment(parameters: dict):
+ """
+ Sets several environment variables in the container to emulate an MWAA environment provisioned
+ within SageMaker Unified Studio. When running in the ECSExecutor, this is a no-op.
+ """
+ import os
+
+ for key, value in parameters.items():
+ os.environ[key] = value
+
+
+with DAG(
+ DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ test_env_id = test_context[ENV_ID_KEY]
+ domain_id = test_context[DOMAIN_ID_KEY]
+ project_id = test_context[PROJECT_ID_KEY]
+ s3_path = test_context[S3_PATH_KEY]
+ region_name = test_context[REGION_NAME_KEY]
+
+ mock_mwaa_environment_params = get_mwaa_environment_params(s3_path)
+
+ setup_mwaa_environment = mock_mwaa_environment(mock_mwaa_environment_params)
Review Comment:
Fixed
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "ruijiang-rjian (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on PR #62962:
URL: https://github.com/apache/airflow/pull/62962#issuecomment-4050618070
Tightening the sagemaker-studio version to make sure it supports receiving domain id and project id parameters configured through SageMakerNotebookOperator
cc: @uranusjr
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on PR #62962:
URL: https://github.com/apache/airflow/pull/62962#issuecomment-4026584430
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
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2927783930
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
+2. A SageMaker Unified Studio Domain (with default VPC and roles)
+3. A project within the SageMaker Unified Studio Domain
+4. A notebook (test_notebook.ipynb) placed in the project's s3 path
+
+This test is identical to example_sagemaker_unified_studio, but passes domain_id, project_id,
+and domain_region explicitly as operator parameters instead of relying on environment variables.
+"""
+
+DAG_ID = "example_sagemaker_unified_studio_explicit_params"
+
+# Externally fetched variables:
+DOMAIN_ID_KEY = "DOMAIN_ID"
+PROJECT_ID_KEY = "PROJECT_ID"
+S3_PATH_KEY = "S3_PATH"
+REGION_NAME_KEY = "REGION_NAME"
+
+sys_test_context_task = (
+ SystemTestContextBuilder()
+ .add_variable(DOMAIN_ID_KEY)
+ .add_variable(PROJECT_ID_KEY)
+ .add_variable(S3_PATH_KEY)
+ .add_variable(REGION_NAME_KEY)
+ .build()
+)
+
+
+def get_mwaa_environment_params(s3_path: str):
+ AIRFLOW_PREFIX = "AIRFLOW__WORKFLOWS__"
+ return {f"{AIRFLOW_PREFIX}PROJECT_S3_PATH": s3_path}
Review Comment:
Fixed
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2897894877
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -78,10 +75,6 @@ def get_mwaa_environment_params(
parameters = {}
parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_ID"] = domain_id
parameters[f"{AIRFLOW_PREFIX}DATAZONE_PROJECT_ID"] = project_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENVIRONMENT_ID"] = environment_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_SCOPE_NAME"] = "dev"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_STAGE"] = "prod"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENDPOINT"] = f"https://datazone.{region_name}.api.aws"
Review Comment:
> Older versions of SDK did use these parameters. But this system test is using the latest version of SDK Airflow can get, so they will just be noops.
We regularly run the system tests against with different AMPP versions and Airflow versions, so let's not assume too much.
> For our use-case at least, the older version of SDK can only be used with the older version of Airflow, so I think it is fine to remove the env vars from here.
How are you enforcing that? Did you update the minimum version of that SDK in the pyproject toml file? I think we likely did, but let's double check that.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2897658573
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -78,10 +75,6 @@ def get_mwaa_environment_params(
parameters = {}
parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_ID"] = domain_id
parameters[f"{AIRFLOW_PREFIX}DATAZONE_PROJECT_ID"] = project_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENVIRONMENT_ID"] = environment_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_SCOPE_NAME"] = "dev"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_STAGE"] = "prod"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENDPOINT"] = f"https://datazone.{region_name}.api.aws"
Review Comment:
Older versions of SDK did use these parameters. But this system test is using the latest version of SDK Airflow can get, so they will just be noops.
For our use-case at least, the older version of SDK can only be used with the older version of Airflow, so I think it is fine to remove the env vars from 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2897913633
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py:
##########
@@ -101,6 +101,11 @@ class SageMakerNotebookOperator(BaseOperator):
"""
operator_extra_links = (SageMakerUnifiedStudioLink(),)
+ # These fields are declared as template_fields so Airflow resolves XCom references
+ # (e.g. task_instance.xcom_pull(...)) to actual string values before execute() is called.
+ # Without this, the hook would be instantiated with unresolved PlainXComArg objects,
+ # causing a ParamValidationError when the underlying SDK tries to use them as strings.
+ template_fields = ("domain_id", "project_id", "domain_region")
Review Comment:
This was marked as resolved but no change was made or answer to the question. Unresolving this comment.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2927780459
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py:
##########
@@ -101,6 +101,11 @@ class SageMakerNotebookOperator(BaseOperator):
"""
operator_extra_links = (SageMakerUnifiedStudioLink(),)
+ # These fields are declared as template_fields so Airflow resolves XCom references
+ # (e.g. task_instance.xcom_pull(...)) to actual string values before execute() is called.
+ # Without this, the hook would be instantiated with unresolved PlainXComArg objects,
+ # causing a ParamValidationError when the underlying SDK tries to use them as strings.
+ template_fields = ("domain_id", "project_id", "domain_region")
Review Comment:
Fixed
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2932005801
##########
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:
Good point, fixing in the next rev
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2942859325
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -124,9 +131,35 @@ def mock_mwaa_environment(parameters: dict):
setup_mwaa_environment = mock_mwaa_environment(mock_mwaa_environment_params)
- # [START howto_operator_sagemaker_unified_studio_notebook]
notebook_path = "test_notebook.ipynb" # This should be the path to your .ipynb, .sqlnb, or .vetl file in your project.
+ # [START howto_operator_sagemaker_unified_studio_notebook_explicit_params]
Review Comment:
This is adding 1 more task to test the edge-case when the additional vars are passed as parameters to SageMakerNotebookOperator
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "ruijiang-rjian (via GitHub)" <gi...@apache.org>.
ruijiang-rjian commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2932932643
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py:
##########
@@ -113,7 +113,14 @@ def __init__(
def _get_sagemaker_studio_config(self):
config = ClientConfig()
- config.overrides["execution"] = {"local": is_local_runner()}
+ if self.domain_region:
+ config.region = self.domain_region
Review Comment:
Nice catch!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2932011891
##########
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:
Yes, I can push this to mainline, so that it goes out to latest SDK in the next release
##########
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:
Sure
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2897700099
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
Review Comment:
This is carried over from the existing test. I think that is talking about AWS authentication setup in general.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2892414832
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
Review Comment:
Why are you creating a whole new system test for the same operator? This has a lot of downstream effects, we should try to avoid this if we can. You can just put another task which calls this operator in the existing system test.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py:
##########
@@ -101,6 +101,11 @@ class SageMakerNotebookOperator(BaseOperator):
"""
operator_extra_links = (SageMakerUnifiedStudioLink(),)
+ # These fields are declared as template_fields so Airflow resolves XCom references
+ # (e.g. task_instance.xcom_pull(...)) to actual string values before execute() is called.
+ # Without this, the hook would be instantiated with unresolved PlainXComArg objects,
+ # causing a ParamValidationError when the underlying SDK tries to use them as strings.
+ template_fields = ("domain_id", "project_id", "domain_region")
Review Comment:
Any reason to only want to template these ones? What about some of the other inputs to this operator?
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -78,10 +75,6 @@ def get_mwaa_environment_params(
parameters = {}
parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_ID"] = domain_id
parameters[f"{AIRFLOW_PREFIX}DATAZONE_PROJECT_ID"] = project_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENVIRONMENT_ID"] = environment_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_SCOPE_NAME"] = "dev"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_STAGE"] = "prod"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENDPOINT"] = f"https://datazone.{region_name}.api.aws"
Review Comment:
Is it helpful to keep these to test the path for folks using older versions of the SMUS SDK? Is that a relevant path even?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #62962:
URL: https://github.com/apache/airflow/pull/62962#issuecomment-4043135350
@Shnekit This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **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](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst).
- :x: **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](https://github.com/apache/airflow/blob/main/contributing-docs/12_provider_distributions.rst).
- :warning: **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](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst).
> **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](https://s.apache.org/airflow-slack).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2927775790
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -78,10 +75,6 @@ def get_mwaa_environment_params(
parameters = {}
parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_ID"] = domain_id
parameters[f"{AIRFLOW_PREFIX}DATAZONE_PROJECT_ID"] = project_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENVIRONMENT_ID"] = environment_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_SCOPE_NAME"] = "dev"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_STAGE"] = "prod"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENDPOINT"] = f"https://datazone.{region_name}.api.aws"
Review Comment:
New rev with the legacy test as well as SDK version bump
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2942514718
##########
providers/amazon/pyproject.toml:
##########
@@ -76,7 +76,7 @@ dependencies = [
"PyAthena>=3.10.0",
"jmespath>=0.7.0",
# Do not change sagemaker-studio versions without approval from AWS.
- "sagemaker-studio>=1.0.9",
+ "sagemaker-studio>=1.0.25,<1.1.0",
# Sagemaker studio in Python 3.13 requires version >=1.1.0 and Pydynamodb >=0.7.5
"pydynamodb>=0.8.1; python_version >= '3.13'",
Review Comment:
This should be removed as no longer needed
##########
providers/amazon/pyproject.toml:
##########
@@ -76,7 +76,7 @@ dependencies = [
"PyAthena>=3.10.0",
"jmespath>=0.7.0",
# Do not change sagemaker-studio versions without approval from AWS.
- "sagemaker-studio>=1.0.9",
+ "sagemaker-studio>=1.0.25,<1.1.0",
Review Comment:
There should be a comment explaining the reasoning for the limit of <1.1.0 and explnation how the airflow can track it's progress given that the upstream library is not open sourced.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py:
##########
@@ -101,6 +101,20 @@ class SageMakerNotebookOperator(BaseOperator):
"""
operator_extra_links = (SageMakerUnifiedStudioLink(),)
+ # These fields are declared as template_fields so Airflow resolves XCom references
+ # (e.g. task_instance.xcom_pull(...)) to actual string values before execute() is called.
+ # Without this, the hook would be instantiated with unresolved PlainXComArg objects,
+ # causing a ParamValidationError when the underlying SDK tries to use them as strings.
Review Comment:
There is no need for this comment. This is known to airflow developers
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -124,9 +131,35 @@ def mock_mwaa_environment(parameters: dict):
setup_mwaa_environment = mock_mwaa_environment(mock_mwaa_environment_params)
- # [START howto_operator_sagemaker_unified_studio_notebook]
notebook_path = "test_notebook.ipynb" # This should be the path to your .ipynb, .sqlnb, or .vetl file in your project.
+ # [START howto_operator_sagemaker_unified_studio_notebook_explicit_params]
Review Comment:
The `howto_operator_sagemaker_unified_studio_notebook_explicit_params` marker is not refrenced anywhere in the docs. Why?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2893344371
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -78,10 +75,6 @@ def get_mwaa_environment_params(
parameters = {}
parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_ID"] = domain_id
parameters[f"{AIRFLOW_PREFIX}DATAZONE_PROJECT_ID"] = project_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENVIRONMENT_ID"] = environment_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_SCOPE_NAME"] = "dev"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_STAGE"] = "prod"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENDPOINT"] = f"https://datazone.{region_name}.api.aws"
Review Comment:
> these are not being actually used in SDK either so we completely removed these in SMUS SDK recently
I'm getting mixed signals. You removed them recently, so in the past they were being used, right? Or is it that they were never ever used in the SDK?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "ruijiang-rjian (via GitHub)" <gi...@apache.org>.
ruijiang-rjian commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2892706891
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
Review Comment:
sry what's IAM IDC.. do you mean IDC([IAM Identity Center](https://aws.amazon.com/iam/identity-center/))?
Or IAM domain or IDC domain...?
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
+2. A SageMaker Unified Studio Domain (with default VPC and roles)
+3. A project within the SageMaker Unified Studio Domain
+4. A notebook (test_notebook.ipynb) placed in the project's s3 path
+
+This test is identical to example_sagemaker_unified_studio, but passes domain_id, project_id,
+and domain_region explicitly as operator parameters instead of relying on environment variables.
+"""
+
+DAG_ID = "example_sagemaker_unified_studio_explicit_params"
+
+# Externally fetched variables:
+DOMAIN_ID_KEY = "DOMAIN_ID"
+PROJECT_ID_KEY = "PROJECT_ID"
+S3_PATH_KEY = "S3_PATH"
+REGION_NAME_KEY = "REGION_NAME"
+
+sys_test_context_task = (
+ SystemTestContextBuilder()
+ .add_variable(DOMAIN_ID_KEY)
+ .add_variable(PROJECT_ID_KEY)
+ .add_variable(S3_PATH_KEY)
+ .add_variable(REGION_NAME_KEY)
+ .build()
+)
+
+
+def get_mwaa_environment_params(s3_path: str):
+ AIRFLOW_PREFIX = "AIRFLOW__WORKFLOWS__"
+ return {f"{AIRFLOW_PREFIX}PROJECT_S3_PATH": s3_path}
Review Comment:
no we shouldn't need to use this(and actually our system test should prove that we don't need any env variables to be set in the airflow env for the operator with new parameters provided as inputs to work)
we're expecting to be able to infer the s3 path using the provided domain_id and project_id: https://code.amazon.com/packages/MaxDomePythonSDK/blobs/947dbe39c55874437a8b9c5e181a3d3483c85ca9/--/src/sagemaker_studio/execution/remote_execution_client.py#L887,L888
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
+2. A SageMaker Unified Studio Domain (with default VPC and roles)
+3. A project within the SageMaker Unified Studio Domain
+4. A notebook (test_notebook.ipynb) placed in the project's s3 path
+
+This test is identical to example_sagemaker_unified_studio, but passes domain_id, project_id,
+and domain_region explicitly as operator parameters instead of relying on environment variables.
+"""
+
+DAG_ID = "example_sagemaker_unified_studio_explicit_params"
+
+# Externally fetched variables:
+DOMAIN_ID_KEY = "DOMAIN_ID"
+PROJECT_ID_KEY = "PROJECT_ID"
+S3_PATH_KEY = "S3_PATH"
+REGION_NAME_KEY = "REGION_NAME"
+
+sys_test_context_task = (
+ SystemTestContextBuilder()
+ .add_variable(DOMAIN_ID_KEY)
+ .add_variable(PROJECT_ID_KEY)
+ .add_variable(S3_PATH_KEY)
+ .add_variable(REGION_NAME_KEY)
+ .build()
+)
+
+
+def get_mwaa_environment_params(s3_path: str):
+ AIRFLOW_PREFIX = "AIRFLOW__WORKFLOWS__"
+ return {f"{AIRFLOW_PREFIX}PROJECT_S3_PATH": s3_path}
+
+
+@task
+def mock_mwaa_environment(parameters: dict):
+ """
+ Sets several environment variables in the container to emulate an MWAA environment provisioned
+ within SageMaker Unified Studio. When running in the ECSExecutor, this is a no-op.
+ """
+ import os
+
+ for key, value in parameters.items():
+ os.environ[key] = value
+
+
+with DAG(
+ DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ test_env_id = test_context[ENV_ID_KEY]
+ domain_id = test_context[DOMAIN_ID_KEY]
+ project_id = test_context[PROJECT_ID_KEY]
+ s3_path = test_context[S3_PATH_KEY]
+ region_name = test_context[REGION_NAME_KEY]
+
+ mock_mwaa_environment_params = get_mwaa_environment_params(s3_path)
+
+ setup_mwaa_environment = mock_mwaa_environment(mock_mwaa_environment_params)
Review Comment:
we should probably remove the whole mock_mwaa_environment()
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
+2. A SageMaker Unified Studio Domain (with default VPC and roles)
+3. A project within the SageMaker Unified Studio Domain
+4. A notebook (test_notebook.ipynb) placed in the project's s3 path
+
+This test is identical to example_sagemaker_unified_studio, but passes domain_id, project_id,
+and domain_region explicitly as operator parameters instead of relying on environment variables.
+"""
+
+DAG_ID = "example_sagemaker_unified_studio_explicit_params"
+
+# Externally fetched variables:
+DOMAIN_ID_KEY = "DOMAIN_ID"
+PROJECT_ID_KEY = "PROJECT_ID"
+S3_PATH_KEY = "S3_PATH"
+REGION_NAME_KEY = "REGION_NAME"
+
+sys_test_context_task = (
+ SystemTestContextBuilder()
+ .add_variable(DOMAIN_ID_KEY)
+ .add_variable(PROJECT_ID_KEY)
+ .add_variable(S3_PATH_KEY)
+ .add_variable(REGION_NAME_KEY)
+ .build()
+)
+
+
+def get_mwaa_environment_params(s3_path: str):
+ AIRFLOW_PREFIX = "AIRFLOW__WORKFLOWS__"
+ return {f"{AIRFLOW_PREFIX}PROJECT_S3_PATH": s3_path}
+
+
+@task
+def mock_mwaa_environment(parameters: dict):
+ """
+ Sets several environment variables in the container to emulate an MWAA environment provisioned
+ within SageMaker Unified Studio. When running in the ECSExecutor, this is a no-op.
+ """
+ import os
+
+ for key, value in parameters.items():
+ os.environ[key] = value
+
+
+with DAG(
+ DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ test_env_id = test_context[ENV_ID_KEY]
+ domain_id = test_context[DOMAIN_ID_KEY]
+ project_id = test_context[PROJECT_ID_KEY]
+ s3_path = test_context[S3_PATH_KEY]
+ region_name = test_context[REGION_NAME_KEY]
+
+ mock_mwaa_environment_params = get_mwaa_environment_params(s3_path)
+
+ setup_mwaa_environment = mock_mwaa_environment(mock_mwaa_environment_params)
+
+ # [START howto_operator_sagemaker_unified_studio_notebook_explicit_params]
+ 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="run-notebook",
+ domain_id=domain_id,
+ project_id=project_id,
+ domain_region=region_name,
Review Comment:
if possible, we'd also need to test without providing domain_region as it's an optional field..
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py:
##########
@@ -101,6 +101,11 @@ class SageMakerNotebookOperator(BaseOperator):
"""
operator_extra_links = (SageMakerUnifiedStudioLink(),)
+ # These fields are declared as template_fields so Airflow resolves XCom references
+ # (e.g. task_instance.xcom_pull(...)) to actual string values before execute() is called.
+ # Without this, the hook would be instantiated with unresolved PlainXComArg objects,
+ # causing a ParamValidationError when the underlying SDK tries to use them as strings.
+ template_fields = ("domain_id", "project_id", "domain_region")
Review Comment:
Good catch on this!
and +1 to question above, do we want to apply to other parameters as well?
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio.py:
##########
@@ -78,10 +75,6 @@ def get_mwaa_environment_params(
parameters = {}
parameters[f"{AIRFLOW_PREFIX}DATAZONE_DOMAIN_ID"] = domain_id
parameters[f"{AIRFLOW_PREFIX}DATAZONE_PROJECT_ID"] = project_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENVIRONMENT_ID"] = environment_id
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_SCOPE_NAME"] = "dev"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_STAGE"] = "prod"
- parameters[f"{AIRFLOW_PREFIX}DATAZONE_ENDPOINT"] = f"https://datazone.{region_name}.api.aws"
Review Comment:
these are not being actually used in SDK either so we completely removed these in SMUS SDK recently
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
Review Comment:
yeah I agree, we should try to keep in the same system test as it's the same operator but with a different setup, +1 a new task before the line to set the env variables in MWAA should hopefully work..?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2897911075
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
Review Comment:
Why was this marked as resolved if the duplicate system test is still existing in this diff?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #62962:
URL: https://github.com/apache/airflow/pull/62962#issuecomment-4028474594
Is the version upgrade needed? It’s not explained above if this uses new features.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix SageMakerNotebookOperator domain_id/project_id/domain_region params [airflow]
Posted by "Shnekit (via GitHub)" <gi...@apache.org>.
Shnekit commented on code in PR #62962:
URL: https://github.com/apache/airflow/pull/62962#discussion_r2927786386
##########
providers/amazon/tests/system/amazon/aws/example_sagemaker_unified_studio_explicit_params.py:
##########
@@ -0,0 +1,153 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.sagemaker_unified_studio import (
+ SageMakerNotebookOperator,
+)
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import DAG, chain, task
+else:
+ # Airflow 2 path
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
+ from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+"""
+Prerequisites: The account which runs this test must manually have the following:
+1. An IAM IDC organization set up in the testing region with a user initialized
+2. A SageMaker Unified Studio Domain (with default VPC and roles)
+3. A project within the SageMaker Unified Studio Domain
+4. A notebook (test_notebook.ipynb) placed in the project's s3 path
+
+This test is identical to example_sagemaker_unified_studio, but passes domain_id, project_id,
+and domain_region explicitly as operator parameters instead of relying on environment variables.
+"""
+
+DAG_ID = "example_sagemaker_unified_studio_explicit_params"
+
+# Externally fetched variables:
+DOMAIN_ID_KEY = "DOMAIN_ID"
+PROJECT_ID_KEY = "PROJECT_ID"
+S3_PATH_KEY = "S3_PATH"
+REGION_NAME_KEY = "REGION_NAME"
+
+sys_test_context_task = (
+ SystemTestContextBuilder()
+ .add_variable(DOMAIN_ID_KEY)
+ .add_variable(PROJECT_ID_KEY)
+ .add_variable(S3_PATH_KEY)
+ .add_variable(REGION_NAME_KEY)
+ .build()
+)
+
+
+def get_mwaa_environment_params(s3_path: str):
+ AIRFLOW_PREFIX = "AIRFLOW__WORKFLOWS__"
+ return {f"{AIRFLOW_PREFIX}PROJECT_S3_PATH": s3_path}
+
+
+@task
+def mock_mwaa_environment(parameters: dict):
+ """
+ Sets several environment variables in the container to emulate an MWAA environment provisioned
+ within SageMaker Unified Studio. When running in the ECSExecutor, this is a no-op.
+ """
+ import os
+
+ for key, value in parameters.items():
+ os.environ[key] = value
+
+
+with DAG(
+ DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ test_env_id = test_context[ENV_ID_KEY]
+ domain_id = test_context[DOMAIN_ID_KEY]
+ project_id = test_context[PROJECT_ID_KEY]
+ s3_path = test_context[S3_PATH_KEY]
+ region_name = test_context[REGION_NAME_KEY]
+
+ mock_mwaa_environment_params = get_mwaa_environment_params(s3_path)
+
+ setup_mwaa_environment = mock_mwaa_environment(mock_mwaa_environment_params)
+
+ # [START howto_operator_sagemaker_unified_studio_notebook_explicit_params]
+ 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="run-notebook",
+ domain_id=domain_id,
+ project_id=project_id,
+ domain_region=region_name,
Review Comment:
Added a test for that
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org