Add SageMaker Unified Studio domain_id, project_id, domain_region as new parameters to SageMakerNotebookOperator

docs/spelling_wordlist.txt

@@ -541,6 +541,7 @@ DisplayVideo
 distcp
 distro
 distros
+dkr
 Dlp
 dlp
 DlpJob
@@ -2096,6 +2097,7 @@ xcomresult
 XComs
 Xero
 Xiaodong
+xlarge
 xml
 xmltodict
 xpath

providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py

@@ -40,24 +40,43 @@ class SageMakerNotebookHook(BaseHook):
         from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import SageMakerNotebookHook
 
         notebook_hook = SageMakerNotebookHook(
+            execution_name="notebook_execution",
+            domain_id="dzd-example123456",
+            project_id="example123456",
             input_config={"input_path": "path/to/notebook.ipynb", "input_params": {"param1": "value1"}},
             output_config={"output_uri": "folder/output/location/prefix", "output_formats": "NOTEBOOK"},
-            execution_name="notebook_execution",
+            domain_region="us-east-1",
             waiter_delay=10,
             waiter_max_attempts=1440,
         )
 
     :param execution_name: The name of the notebook job to be executed, this is same as task_id.
+    :param domain_id: The domain ID for Amazon SageMaker Unified Studio. Optional - if not provided,
+        the SDK will attempt to resolve it from the environment.
+    :param project_id: The project ID for Amazon SageMaker Unified Studio. Optional - if not provided,
+        the SDK will attempt to resolve it from the environment.
     :param input_config: Configuration for the input file.
         Example: {'input_path': 'folder/input/notebook.ipynb', 'input_params': {'param1': 'value1'}}
     :param output_config: Configuration for the output format. It should include an output_formats parameter to specify the output format.
         Example: {'output_formats': ['NOTEBOOK']}
-    :param compute: compute configuration to use for the notebook execution. This is a required attribute if the execution is on a remote compute.
-        Example: {"instance_type": "ml.m5.large", "volume_size_in_gb": 30, "volume_kms_key_id": "", "image_details": {"ecr_uri": "string"}, "container_entrypoint": ["string"]}
+    :param domain_region: The AWS region for the domain. If not provided, the default AWS region will be used.
+    :param compute: compute configuration to use for the notebook execution. This is a required attribute
+        if the execution is on a remote compute.
+        Example::
+
+            {
+                "instance_type": "ml.c5.xlarge",
+                "image_details": {
+                    "image_name": "sagemaker-distribution-prod",
+                    "image_version": "3",
+                    "ecr_uri": "123456123456.dkr.ecr.us-west-2.amazonaws.com/ImageName:latest",
+                },
+            }
+
     :param termination_condition: conditions to match to terminate the remote execution.
-        Example: {"MaxRuntimeInSeconds": 3600}
+        Example: ``{"MaxRuntimeInSeconds": 3600}``
     :param tags: tags to be associated with the remote execution runs.
-        Example: {"md_analytics": "logs"}
+        Example: ``{"md_analytics": "logs"}``
     :param waiter_delay: Interval in seconds to check the task execution status.
     :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
     """
@@ -66,7 +85,10 @@ def __init__(
         self,
         execution_name: str,
         input_config: dict | None = None,
+        domain_id: str | None = None,
+        project_id: str | None = None,
         output_config: dict | None = None,
+        domain_region: str | None = None,
         compute: dict | None = None,
         termination_condition: dict | None = None,
         tags: dict | None = None,
@@ -78,6 +100,9 @@ def __init__(
         super().__init__(*args, **kwargs)
         self._sagemaker_studio = SageMakerStudioAPI(self._get_sagemaker_studio_config())
         self.execution_name = execution_name
+        self.domain_id = domain_id
+        self.project_id = project_id
+        self.domain_region = domain_region
         self.input_config = input_config or {}
         self.output_config = output_config or {"output_formats": ["NOTEBOOK"]}
         self.compute = compute
@@ -114,17 +139,20 @@ def start_notebook_execution(self):
         start_execution_params = {
             "execution_name": self.execution_name,
             "execution_type": "NOTEBOOK",
+            "domain_id": self.domain_id,
+            "project_id": self.project_id,
             "input_config": self._format_start_execution_input_config(),
             "output_config": self._format_start_execution_output_config(),
             "termination_condition": self.termination_condition,
             "tags": self.tags,
         }
+
+        if self.domain_region:
+            start_execution_params["domain_region"] = self.domain_region
+
         if self.compute:
             start_execution_params["compute"] = self.compute
-        else:
-            start_execution_params["compute"] = {"instance_type": "ml.m6i.xlarge"}
 
-        print(start_execution_params)
         return self._sagemaker_studio.execution_client.start_execution(**start_execution_params)
 
     def wait_for_execution_completion(self, execution_id, context):

providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py

@@ -48,8 +48,11 @@ class SageMakerNotebookOperator(BaseOperator):
 
         notebook_operator = SageMakerNotebookOperator(
             task_id="notebook_task",
+            domain_id="dzd-example123456",
+            project_id="example123456",
             input_config={"input_path": "path/to/notebook.ipynb", "input_params": ""},
             output_config={"output_format": "ipynb"},
+            domain_region="us-east-1",
             wait_for_completion=True,
             waiter_delay=10,
             waiter_max_attempts=1440,
@@ -63,12 +66,28 @@ class SageMakerNotebookOperator(BaseOperator):
     :param output_config:  Configuration for the output format. It should include an output_format parameter to control
         the format of the notebook execution output.
         Example: {"output_formats": ["NOTEBOOK"]}
-    :param compute: compute configuration to use for the notebook execution. This is a required attribute if the execution is on a remote compute.
-        Example: {"instance_type": "ml.m5.large", "volume_size_in_gb": 30, "volume_kms_key_id": "", "image_details": {"ecr_uri": "string"}, "container_entrypoint": ["string"]}
+    :param domain_id: The domain ID for Amazon SageMaker Unified Studio. Optional - if not provided,
+        the SDK will attempt to resolve it from the environment.
+    :param project_id: The project ID for Amazon SageMaker Unified Studio. Optional - if not provided,
+        the SDK will attempt to resolve it from the environment.
+    :param domain_region: The AWS region for the domain. If not provided, the default AWS region will be used.
+    :param compute: compute configuration to use for the artifact execution. This is a required attribute
+        if the execution is on a remote compute.
+        Example::
+
+            {
+                "instance_type": "ml.c5.xlarge",
+                "image_details": {
+                    "image_name": "sagemaker-distribution-prod",
+                    "image_version": "3",
+                    "ecr_uri": "123456123456.dkr.ecr.us-west-2.amazonaws.com/ImageName:latest",
+                },
+            }
+
     :param termination_condition: conditions to match to terminate the remote execution.
-        Example: { "MaxRuntimeInSeconds": 3600 }
+        Example: ``{"MaxRuntimeInSeconds": 3600}``
     :param tags: tags to be associated with the remote execution runs.
-        Example: { "md_analytics": "logs" }
+        Example: ``{"md_analytics": "logs"}``
     :param wait_for_completion: Indicates whether to wait for the notebook execution to complete. If True, wait for completion; if False, don't wait.
     :param waiter_delay: Interval in seconds to check the notebook execution status.
     :param waiter_max_attempts: Number of attempts to wait before returning FAILED.
@@ -87,7 +106,10 @@ def __init__(
         self,
         task_id: str,
         input_config: dict,
+        domain_id: str | None = None,
+        project_id: str | None = None,
         output_config: dict | None = None,
+        domain_region: str | None = None,
         compute: dict | None = None,
         termination_condition: dict | None = None,
         tags: dict | None = None,
@@ -99,8 +121,11 @@ def __init__(
     ):
         super().__init__(task_id=task_id, **kwargs)
         self.execution_name = task_id
+        self.domain_id = domain_id
+        self.project_id = project_id
         self.input_config = input_config
         self.output_config = output_config or {"output_formats": ["NOTEBOOK"]}
+        self.domain_region = domain_region
         self.compute = compute or {}
         self.termination_condition = termination_condition or {}
         self.tags = tags or {}
@@ -119,9 +144,12 @@ def notebook_execution_hook(self):
             raise AirflowException("input_path is a required field in the input_config")
 
         return SageMakerNotebookHook(
+            domain_id=self.domain_id,
+            project_id=self.project_id,
             input_config=self.input_config,
             output_config=self.output_config,
             execution_name=self.execution_name,
+            domain_region=self.domain_region,
             compute=self.compute,
             termination_condition=self.termination_condition,
             tags=self.tags,

providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker_unified_studio.py

@@ -200,28 +200,6 @@ def test_set_xcom_s3_path_negative_missing_context(self):
         with pytest.raises(AirflowException, match="context is required"):
             self.hook._set_xcom_s3_path(self.s3Path, {})
 
-    def test_start_notebook_execution_default_compute(self):
-        """Test that default compute uses ml.m6i.xlarge instance type."""
-        hook_without_compute = SageMakerNotebookHook(
-            input_config={
-                "input_path": "test-data/notebook/test_notebook.ipynb",
-                "input_params": {"key": "value"},
-            },
-            output_config={"output_formats": ["NOTEBOOK"]},
-            execution_name="test-execution",
-            waiter_delay=10,
-        )
-        hook_without_compute._sagemaker_studio = MagicMock()
-        hook_without_compute._sagemaker_studio.execution_client = MagicMock(spec=ExecutionClient)
-        hook_without_compute._sagemaker_studio.execution_client.start_execution.return_value = {
-            "executionId": "123456"
-        }
-
-        hook_without_compute.start_notebook_execution()
-
-        call_kwargs = hook_without_compute._sagemaker_studio.execution_client.start_execution.call_args[1]
-        assert call_kwargs["compute"] == {"instance_type": "ml.m6i.xlarge"}
-
     def test_start_notebook_execution_custom_compute(self):
         """Test that custom compute config is used when provided."""
         custom_compute = {"instance_type": "ml.c5.xlarge", "volume_size_in_gb": 50}
@@ -245,3 +223,79 @@ def test_start_notebook_execution_custom_compute(self):
 
         call_kwargs = hook_with_compute._sagemaker_studio.execution_client.start_execution.call_args[1]
         assert call_kwargs["compute"] == custom_compute
+
+    def test_init_with_domain_id_project_id_domain_region(self):
+        """Test that domain_id, project_id, and domain_region are stored on the hook."""
+        hook = SageMakerNotebookHook(
+            execution_name="test-execution",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+            domain_id="dzd-example123456",
+            project_id="proj-example123456",
+            domain_region="us-east-1",
+        )
+        hook._sagemaker_studio = MagicMock()
+
+        assert hook.domain_id == "dzd-example123456"
+        assert hook.project_id == "proj-example123456"
+        assert hook.domain_region == "us-east-1"
+
+    def test_init_domain_params_default_to_none(self):
+        """When domain params are not provided they default to None so the SDK can resolve them."""
+        hook = SageMakerNotebookHook(
+            execution_name="test-execution",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+        )
+        hook._sagemaker_studio = MagicMock()
+
+        assert hook.domain_id is None
+        assert hook.project_id is None
+        assert hook.domain_region is None
+
+    def test_start_notebook_execution_includes_domain_id_and_project_id(self):
+        """domain_id and project_id are always forwarded to start_execution."""
+        hook = SageMakerNotebookHook(
+            execution_name="test-execution",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+            domain_id="dzd-example123456",
+            project_id="proj-example123456",
+        )
+        hook._sagemaker_studio = MagicMock()
+        hook._sagemaker_studio.execution_client = MagicMock(spec=ExecutionClient)
+        hook._sagemaker_studio.execution_client.start_execution.return_value = {"execution_id": "abc"}
+
+        hook.start_notebook_execution()
+
+        call_kwargs = hook._sagemaker_studio.execution_client.start_execution.call_args[1]
+        assert call_kwargs["domain_id"] == "dzd-example123456"
+        assert call_kwargs["project_id"] == "proj-example123456"
+
+    def test_start_notebook_execution_includes_domain_region_when_provided(self):
+        """domain_region is conditionally added to start_execution params only when set."""
+        hook = SageMakerNotebookHook(
+            execution_name="test-execution",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+            domain_region="eu-west-1",
+        )
+        hook._sagemaker_studio = MagicMock()
+        hook._sagemaker_studio.execution_client = MagicMock(spec=ExecutionClient)
+        hook._sagemaker_studio.execution_client.start_execution.return_value = {"execution_id": "abc"}
+
+        hook.start_notebook_execution()
+
+        call_kwargs = hook._sagemaker_studio.execution_client.start_execution.call_args[1]
+        assert call_kwargs["domain_region"] == "eu-west-1"
+
+    def test_start_notebook_execution_omits_domain_region_when_not_provided(self):
+        """domain_region must NOT appear in start_execution params when it is None."""
+        hook = SageMakerNotebookHook(
+            execution_name="test-execution",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+        )
+        hook._sagemaker_studio = MagicMock()
+        hook._sagemaker_studio.execution_client = MagicMock(spec=ExecutionClient)
+        hook._sagemaker_studio.execution_client.start_execution.return_value = {"execution_id": "abc"}
+
+        hook.start_notebook_execution()
+
+        call_kwargs = hook._sagemaker_studio.execution_client.start_execution.call_args[1]
+        assert "domain_region" not in call_kwargs

providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio.py

@@ -45,6 +45,29 @@ def test_init(self):
         }
         assert operator.output_config == {"output_format": "ipynb"}
 
+    def test_init_with_domain_id_project_id_domain_region(self):
+        operator = SageMakerNotebookOperator(
+            task_id="test_id",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+            domain_id="dzd-example123456",
+            project_id="proj-example123456",
+            domain_region="us-east-1",
+        )
+
+        assert operator.domain_id == "dzd-example123456"
+        assert operator.project_id == "proj-example123456"
+        assert operator.domain_region == "us-east-1"
+
+    def test_init_domain_params_default_to_none(self):
+        operator = SageMakerNotebookOperator(
+            task_id="test_id",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+        )
+
+        assert operator.domain_id is None
+        assert operator.project_id is None
+        assert operator.domain_region is None
+
     def test_only_required_params_init(self):
         operator = SageMakerNotebookOperator(
             task_id="test_id",
@@ -54,6 +77,59 @@ def test_only_required_params_init(self):
         )
         assert isinstance(operator, SageMakerNotebookOperator)
 
+    @patch("airflow.providers.amazon.aws.operators.sagemaker_unified_studio.SageMakerNotebookHook")
+    def test_execute_passes_domain_id_project_id_domain_region_to_hook(self, mock_notebook_hook):
+        mock_hook_instance = mock_notebook_hook.return_value
+        mock_hook_instance.start_notebook_execution.return_value = {
+            "execution_id": "123456",
+            "executionType": "test",
+        }
+
+        operator = SageMakerNotebookOperator(
+            task_id="test_id",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+            domain_id="dzd-example123456",
+            project_id="proj-example123456",
+            domain_region="us-west-2",
+        )
+
+        operator.execute({})
+
+        mock_notebook_hook.assert_called_once_with(
+            domain_id="dzd-example123456",
+            project_id="proj-example123456",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+            output_config={"output_formats": ["NOTEBOOK"]},
+            execution_name="test_id",
+            domain_region="us-west-2",
+            compute={},
+            termination_condition={},
+            tags={},
+            waiter_delay=10,
+            waiter_max_attempts=1440,
+        )
+
+    @patch("airflow.providers.amazon.aws.operators.sagemaker_unified_studio.SageMakerNotebookHook")
+    def test_execute_passes_none_domain_params_to_hook(self, mock_notebook_hook):
+        """When domain_id/project_id/domain_region are omitted, None is forwarded so the SDK resolves them."""
+        mock_hook_instance = mock_notebook_hook.return_value
+        mock_hook_instance.start_notebook_execution.return_value = {
+            "execution_id": "123456",
+            "executionType": "test",
+        }
+
+        operator = SageMakerNotebookOperator(
+            task_id="test_id",
+            input_config={"input_path": "path/to/notebook.ipynb"},
+        )
+
+        operator.execute({})
+
+        call_kwargs = mock_notebook_hook.call_args[1]
+        assert call_kwargs["domain_id"] is None
+        assert call_kwargs["project_id"] is None
+        assert call_kwargs["domain_region"] is None
+
     @patch("airflow.providers.amazon.aws.operators.sagemaker_unified_studio.SageMakerNotebookHook")
     def test_execute_success(self, mock_notebook_hook):  # Mock the NotebookHook and its execute method
         mock_hook_instance = mock_notebook_hook.return_value