Repository: apache/airflow
File: providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker_unified_studio.py (Python)

Code before:

class SageMakerNotebookHook(BaseHook):
    """
    Interact with Sagemaker Unified Studio Workflows.

    This hook provides a wrapper around the Sagemaker Workflows Notebook Execution API.

    Examples:
     .. code-block:: python

        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"},
            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 MWAA 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 MWAA 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 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}``
    :param tags: tags to be associated with the remote execution runs.
        Example: ``{"md_analytics": "logs"}``
    :param waiter_delay: Interval in seconds to check the task execution status.

Review comment (suggestion):
Carrying over this feedback from the other PR that doesn't look like it was addressed:


It is technically possible to export the right variables in a non-MWAA env, no?

```suggestion
        the SDK will attempt to resolve it from the environment.
```


Same suggestion for below and in the operator module.

Code after:

class SageMakerNotebookHook(BaseHook):
    """
    Interact with Sagemaker Unified Studio Workflows.

    This hook provides a wrapper around the Sagemaker Workflows Notebook Execution API.

    Examples:
     .. code-block:: python

        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"},
            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 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}``
    :param tags: tags to be associated with the remote execution runs.
        Example: ``{"md_analytics": "logs"}``
    :param waiter_delay: Interval in seconds to check the task execution status.