vincbeck commented on code in PR #33219:
URL: https://github.com/apache/airflow/pull/33219#discussion_r1294686733
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
Review Comment:
Why do you need to store it as instance variable? I would move this part in
`execute`, no need to save it in the instance
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.wait_for_completion = wait_for_completion
+ self.aws_conn_id = aws_conn_id
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Stopping SageMaker notebook %s.", self.instance_name)
+
self.hook.conn.stop_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to stop",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_stopped").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerDeleteNotebookOperator(BaseOperator):
+ """
+
+ Delete a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerDeleteNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to delete.
+ :param config: Additional configuration options for the delete call.
+ :param wait_for_completion: Whether or not to wait for the notebook to
delete before returning.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.wait_for_completion = wait_for_completion
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
Review Comment:
```suggestion
def hook(self) -> SageMakerHook:
"""Create and return SageMakerHook."""
return SageMakerHook(aws_conn_id=self.aws_conn_id)
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.wait_for_completion = wait_for_completion
+ self.aws_conn_id = aws_conn_id
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Stopping SageMaker notebook %s.", self.instance_name)
+
self.hook.conn.stop_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to stop",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_stopped").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerDeleteNotebookOperator(BaseOperator):
+ """
+
+ Delete a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerDeleteNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to delete.
+ :param config: Additional configuration options for the delete call.
+ :param wait_for_completion: Whether or not to wait for the notebook to
delete before returning.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
Review Comment:
Not used
```suggestion
```
##########
airflow/providers/amazon/aws/hooks/sagemaker.py:
##########
@@ -1305,3 +1305,23 @@ def create_auto_ml_job(
if "BestCandidate" in res:
return res["BestCandidate"]
return None
+
+
+class SageMakerNotebookHook(AwsBaseHook):
Review Comment:
This is not needed. `SageMakerHook` is enough and does the same thing
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
Review Comment:
```suggestion
def hook(self) -> SageMakerHook:
"""Create and return SageMakerHook."""
return SageMakerHook(aws_conn_id=self.aws_conn_id)
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.wait_for_completion = wait_for_completion
+ self.aws_conn_id = aws_conn_id
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Stopping SageMaker notebook %s.", self.instance_name)
+
self.hook.conn.stop_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to stop",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_stopped").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerDeleteNotebookOperator(BaseOperator):
+ """
+
+ Delete a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerDeleteNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to delete.
+ :param config: Additional configuration options for the delete call.
+ :param wait_for_completion: Whether or not to wait for the notebook to
delete before returning.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.wait_for_completion = wait_for_completion
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Deleting SageMaker notebook %s....", self.instance_name)
+
self.hook.conn.delete_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to delete...",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_deleted").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerStartNoteBookOperator(BaseOperator):
+ """
+
+ Start a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStartNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to start.
+ :param wait_for_completion: Whether or not to wait for notebook to be
InService before returning
+ :param config: Additional configuration options for the start call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
Review Comment:
Not used
```suggestion
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
Review Comment:
`trim_none_values` is deprecated
```suggestion
**prune_dict(self.create_notebook_instance_kwargs)
```
##########
tests/providers/amazon/aws/operators/test_sagemaker_notebook.py:
##########
@@ -0,0 +1,167 @@
+#
+# 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 unittest import mock
+
+import pytest
+from moto import mock_sagemaker
+
+from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerNotebookHook
+from airflow.providers.amazon.aws.operators.sagemaker import (
+ SageMakerCreateNotebookOperator,
+ SageMakerDeleteNotebookOperator,
+ SageMakerStartNoteBookOperator,
+ SageMakerStopNotebookOperator,
+)
+
+INSTANCE_NAME = "notebook"
+INSTANCE_TYPE = "ml.t3.medium"
+ROLE_ARN = "arn:aws:iam:role/role"
+
+
+@pytest.fixture
+def hook() -> SageMakerNotebookHook:
+ with mock_sagemaker():
+ yield SageMakerNotebookHook(aws_conn_id="aws_default")
+
+
+@pytest.fixture
+def create_instance_args():
+ return {
+ "NotebookInstanceName": INSTANCE_NAME,
+ "InstanceType": INSTANCE_TYPE,
+ "RoleArn": ROLE_ARN,
+ }
+
+
+class TestSageMakerNotebookHook:
+ def test_conn(self):
+ hook = SageMakerNotebookHook(aws_conn_id="sagemaker_test_conn_id")
+ assert hook.aws_conn_id == "sagemaker_test_conn_id"
+
+ def test_create_instance(self, hook: SageMakerNotebookHook,
create_instance_args, capsys):
+ # create a notebook
+ resp = hook.conn.create_notebook_instance(**create_instance_args)
+ assert resp["NotebookInstanceArn"]
+
+ def test_start_instance(self, hook, create_instance_args, capsys):
+ hook.conn.create_notebook_instance(**create_instance_args)
+ resp =
hook.conn.start_notebook_instance(NotebookInstanceName=INSTANCE_NAME)
+ assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+ def test_stop_instance(self, hook, create_instance_args, capsys):
+ hook.conn.create_notebook_instance(**create_instance_args)
+ resp =
hook.conn.stop_notebook_instance(NotebookInstanceName=INSTANCE_NAME)
+ assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+ def test_delete_instance(self, hook, create_instance_args, capsys):
+ hook.conn.create_notebook_instance(**create_instance_args)
+ hook.conn.stop_notebook_instance(NotebookInstanceName=INSTANCE_NAME)
+ resp =
hook.conn.delete_notebook_instance(NotebookInstanceName=INSTANCE_NAME)
+ assert resp["ResponseMetadata"]["HTTPStatusCode"] == 200
+
+
+class TestSagemakerCreateNotebookOperator:
+ @mock.patch.object(SageMakerNotebookHook, "conn")
+ def test_create_notebook_without_wait_for_completion(self, mock_hook_conn):
+ operator = SageMakerCreateNotebookOperator(
+ task_id="task_test",
+ instance_name=INSTANCE_NAME,
+ instance_type=INSTANCE_TYPE,
+ role_arn=ROLE_ARN,
+ wait_for_completion=False,
+ volume_size_in_gb=50,
+ )
+ operator.execute(None)
+ mock_hook_conn.create_notebook_instance.assert_called_once()
+ mock_hook_conn.get_waiter.assert_not_called()
+
+ # test wait_for_completion
Review Comment:
```suggestion
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
Review Comment:
Not used, you can remove it
```suggestion
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.wait_for_completion = wait_for_completion
+ self.aws_conn_id = aws_conn_id
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Stopping SageMaker notebook %s.", self.instance_name)
+
self.hook.conn.stop_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to stop",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_stopped").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerDeleteNotebookOperator(BaseOperator):
+ """
+
+ Delete a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerDeleteNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to delete.
+ :param config: Additional configuration options for the delete call.
+ :param wait_for_completion: Whether or not to wait for the notebook to
delete before returning.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.wait_for_completion = wait_for_completion
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Deleting SageMaker notebook %s....", self.instance_name)
+
self.hook.conn.delete_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to delete...",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_deleted").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerStartNoteBookOperator(BaseOperator):
+ """
+
+ Start a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStartNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to start.
+ :param wait_for_completion: Whether or not to wait for notebook to be
InService before returning
+ :param config: Additional configuration options for the start call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.wait_for_completion = wait_for_completion
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
Review Comment:
```suggestion
def hook(self) -> SageMakerHook:
"""Create and return SageMakerHook."""
return SageMakerHook(aws_conn_id=self.aws_conn_id)
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.wait_for_completion = wait_for_completion
+ self.aws_conn_id = aws_conn_id
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
Review Comment:
```suggestion
def hook(self) -> SageMakerHook:
"""Create and return SageMakerHook."""
return SageMakerHook(aws_conn_id=self.aws_conn_id)
```
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.wait_for_completion = wait_for_completion
+ self.aws_conn_id = aws_conn_id
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Stopping SageMaker notebook %s.", self.instance_name)
+
self.hook.conn.stop_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to stop",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_stopped").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerDeleteNotebookOperator(BaseOperator):
+ """
+
Review Comment:
```suggestion
```
##########
tests/system/providers/amazon/aws/example_sagemaker_notebook.py:
##########
@@ -0,0 +1,109 @@
+# 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 import DAG
+from airflow.models.baseoperator import chain
+from airflow.providers.amazon.aws.operators.sagemaker import (
+ SageMakerCreateNotebookOperator,
+ SageMakerDeleteNotebookOperator,
+ SageMakerStartNoteBookOperator,
+ SageMakerStopNotebookOperator,
+)
+from tests.system.providers.amazon.aws.utils import ENV_ID_KEY,
SystemTestContextBuilder
+
+DAG_ID = "example_sagemaker_notebook"
+
+# Externally fetched variables:
+ROLE_ARN_KEY = "ROLE_ARN" # must have an IAM role to run notebooks
+
+sys_test_context_task =
SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
+
+with DAG(
+ DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ tags=["example"],
+ catchup=False,
+) as dag:
+
+ test_context = sys_test_context_task()
+
+ instance_name: str = f"{test_context[ENV_ID_KEY]}-test-notebook"
+
+ role_arn = test_context[ROLE_ARN_KEY]
+
+ # [START howto_operator_sagemaker_notebook_create]
+ instance = SageMakerCreateNotebookOperator(
+ task_id="create_instance",
+ instance_name=instance_name,
+ instance_type="ml.t3.medium",
+ role_arn=role_arn,
+ wait_for_completion=True,
+ )
+ # [END howto_operator_sagemaker_notebook_create]
+
+ # [START howto_operator_sagemaker_notebook_stop]
+ stop_instance = SageMakerStopNotebookOperator(
+ task_id="stop_instance",
+ instance_name=instance_name,
+ )
+ # [END howto_operator_sagemaker_notebook_stop]
+
+ # [START howto_operator_sagemaker_notebook_start]
+ start_instance = SageMakerStartNoteBookOperator(
+ task_id="start_instance",
+ instance_name=instance_name,
+ )
+
+ # [END howto_operator_sagemaker_notebook_start]
+
+ # [START howto_operator_sagemaker_notebook_delete]
+ # Instance must be stopped before it can be deleted.
+ stop_instance_before_delete = SageMakerStopNotebookOperator(
+ task_id="stop_instance_before_delete",
+ instance_name=instance_name,
+ )
+ delete_instance =
SageMakerDeleteNotebookOperator(task_id="delete_instance",
instance_name=instance_name)
+ # [END howto_operator_sagemaker_notebook_delete]
Review Comment:
No need to include the stop operation in docs in the delete section
```suggestion
# Instance must be stopped before it can be deleted.
stop_instance_before_delete = SageMakerStopNotebookOperator(
task_id="stop_instance_before_delete",
instance_name=instance_name,
)
# [START howto_operator_sagemaker_notebook_delete]
delete_instance =
SageMakerDeleteNotebookOperator(task_id="delete_instance",
instance_name=instance_name)
# [END howto_operator_sagemaker_notebook_delete]
```
##########
scripts/ci/docker-compose/empty/git_version:
##########
Review Comment:
Remove this file from the PR please
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1523,3 +1523,259 @@ def execute(self, context: Context) -> str:
arn = ans["ExperimentArn"]
self.log.info("Experiment %s created successfully with ARN %s.",
self.name, arn)
return arn
+
+
+class SageMakerCreateNotebookOperator(BaseOperator):
+ """
+ Create a SageMaker notebook.
+
+ More information regarding parameters of this operator can be found here
+
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/create_notebook_instance.html.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerCreateNotebookOperator`
+
+ :param instance_name: The name of the notebook instance.
+ :param instance_type: The type of instance to create.
+ :param role_arn: The Amazon Resource Name (ARN) of the IAM role that
SageMaker can assume to access
+ :param volume_size_in_gb: Size in GB of the EBS root device volume of the
notebook instance.
+ :param volume_kms_key_id: The KMS key ID for the EBS root device volume.
+ :param lifecycle_config_name: The name of the lifecycle configuration to
associate with the notebook
+ :param direct_internet_access: Whether to enable direct internet access
for the notebook instance.
+ :param root_access: Whether to give the notebook instance root access to
the Amazon S3 bucket.
+ :param wait_for_completion: Whether or not to wait for the notebook to be
InService before returning
+ :param create_instance_kwargs: Additional configuration options for the
create call.
+ :param config: Additional configuration options for the create call.
+ :param aws_conn_id: The AWS connection ID to use.
+
+ This operator returns The ARN of the created notebook.
+ """
+
+ template_fields: Sequence[str] = (
+ "instance_name",
+ "instance_type",
+ "role_arn",
+ "volume_size_in_gb",
+ "volume_kms_key_id",
+ "lifecycle_config_name",
+ "direct_internet_access",
+ "root_access",
+ "wait_for_completion",
+ "create_instance_kwargs",
+ "config",
+ )
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ *,
+ instance_name: str,
+ instance_type: str,
+ role_arn: str,
+ volume_size_in_gb: int | None = None,
+ volume_kms_key_id: str | None = None,
+ lifecycle_config_name: str | None = None,
+ direct_internet_access: str | None = None,
+ root_access: str | None = None,
+ create_instance_kwargs: dict[str, Any] = {},
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.instance_type = instance_type
+ self.role_arn = role_arn
+ self.volume_size_in_gb = volume_size_in_gb
+ self.volume_kms_key_id = volume_kms_key_id
+ self.lifecycle_config_name = lifecycle_config_name
+ self.direct_internet_access = direct_internet_access
+ self.root_access = root_access
+ self.wait_for_completion = wait_for_completion
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.create_instance_kwargs = create_instance_kwargs
+
+ if "tags" in self.create_instance_kwargs and
self.create_instance_kwargs["tags"] is not None:
+ self.create_instance_kwargs["tags"] =
format_tags(self.create_instance_kwargs["tags"])
+
+ self.create_notebook_instance_kwargs = {
+ "NotebookInstanceName": self.instance_name,
+ "InstanceType": self.instance_type,
+ "RoleArn": self.role_arn,
+ "VolumeSizeInGB": self.volume_size_in_gb,
+ "KmsKeyId": self.volume_kms_key_id,
+ "LifecycleConfigName": self.lifecycle_config_name,
+ "DirectInternetAccess": self.direct_internet_access,
+ "RootAccess": self.root_access,
+ }
+ if len(self.create_instance_kwargs) > 0:
+
self.create_notebook_instance_kwargs.update(self.create_instance_kwargs)
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context: Context):
+ self.log.info("Creating SageMaker notebook %s.", self.instance_name)
+ response = self.hook.conn.create_notebook_instance(
+ **trim_none_values(self.create_notebook_instance_kwargs)
+ )
+
+ self.log.info("SageMaker notebook created: %s",
response["NotebookInstanceArn"])
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to be in
service", self.instance_name)
+ waiter = self.hook.conn.get_waiter("notebook_instance_in_service")
+ waiter.wait(NotebookInstanceName=self.instance_name)
+
+ return response["NotebookInstanceArn"]
+
+
+class SageMakerStopNotebookOperator(BaseOperator):
+ """
+ Stop a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerStopNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to stop.
+ :param config: Additional configuration options for the stop call.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.wait_for_completion = wait_for_completion
+ self.aws_conn_id = aws_conn_id
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Stopping SageMaker notebook %s.", self.instance_name)
+
self.hook.conn.stop_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to stop",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_stopped").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerDeleteNotebookOperator(BaseOperator):
+ """
+
+ Delete a notebook instance.
+
+ .. seealso:
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:SageMakerDeleteNotebookOperator`
+
+ :param instance_name: The name of the notebook instance to delete.
+ :param config: Additional configuration options for the delete call.
+ :param wait_for_completion: Whether or not to wait for the notebook to
delete before returning.
+ :param aws_conn_id: The AWS connection ID to use.
+ """
+
+ template_fields: Sequence[str] = ("instance_name", "wait_for_completion",
"config")
+
+ ui_color = "#ff7300"
+
+ def __init__(
+ self,
+ instance_name: str,
+ wait_for_completion: bool = True,
+ config: dict = {},
+ aws_conn_id: str = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.instance_name = instance_name
+ self.config = config
+ self.aws_conn_id = aws_conn_id
+ self.wait_for_completion = wait_for_completion
+
+ @cached_property
+ def hook(self) -> SageMakerNotebookHook:
+ """Create and return SageMakerNotebookHook."""
+ return SageMakerNotebookHook(aws_conn_id=self.aws_conn_id)
+
+ def execute(self, context):
+ self.log.info("Deleting SageMaker notebook %s....", self.instance_name)
+
self.hook.conn.delete_notebook_instance(NotebookInstanceName=self.instance_name)
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for SageMaker notebook %s to delete...",
self.instance_name)
+ self.hook.conn.get_waiter("notebook_instance_deleted").wait(
+ NotebookInstanceName=self.instance_name
+ )
+
+
+class SageMakerStartNoteBookOperator(BaseOperator):
+ """
+
Review Comment:
```suggestion
```
--
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-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
