diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 56bfd885844cc..205e68fd66a24 100644 --- a/docs/spelling_wordlist.txt +++ b/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 diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py index 0895ffc3f8be2..ac06fcbacf2d9 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py +++ b/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): diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py index 85de6dd1c427c..703108736703c 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker_unified_studio.py +++ b/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, diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker_unified_studio.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker_unified_studio.py index be81da9b2824c..fd27931d22856 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker_unified_studio.py +++ b/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 diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio.py b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio.py index 94542d03a5a89..2f402475d1745 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_unified_studio.py +++ b/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