vincbeck commented on code in PR #33219: URL: https://github.com/apache/airflow/pull/33219#discussion_r1287491411
########## airflow/providers/amazon/aws/hooks/sagemaker_notebook.py: ########## @@ -0,0 +1,153 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import Any + +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook +from airflow.providers.amazon.aws.utils import trim_none_values +from airflow.providers.amazon.aws.utils.tags import format_tags + + +class SageMakerNotebookHook(AwsBaseHook): + """Interact with Amazon SageMaker to execute notebooks. + + Provide thick wrapper around + :external+boto3:py:class:`boto3.client('sagemaker') <SageMaker.Client>` + + Additional arguments (such as ``aws_conn_id``) may be specified and + are passed down to the underlying AwsBaseHook. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, *args, **kwargs) -> None: + kwargs["client_type"] = "sagemaker" + super().__init__(*args, **kwargs) + + def create_instance( + self, + instance_name: str, + instance_type: str, + role_arn: str, + volume_size_in_gb: int | None = None, + volume_kms_key_id: str | None = None, + tags: dict = {}, + subnet_id: str | None = None, + security_group_ids: list = [], + lifecycle_config_name: str | None = None, + direct_internet_access: str | None = None, + accelerator_types: list = [], + default_code_repo: list = [], + additional_code_repos: list = [], + root_access: str | None = None, + platform_id: str | None = None, + imds_config: dict | None = None, + **kwargs: Any, + ) -> dict: + """Create a SageMaker notebook instance. + + :param instance_name: The name of the notebook instance. + :param instance_type: The type of instance to create. For example, 'ml.t2.medium'. + :param image_uri: The Amazon EC2 Image URI for the SageMaker image to use. + :param role_arn: The Amazon Resource Name (ARN) of the IAM role to associate with the notebook + instance. + :param volume_size_in_gb: The size of the EBS volume to attach to the notebook instance. + :param volume_kms_key_id: The KMS key ID to use when creating the notebook instance. + :param tags: A list of tags to associate with the notebook instance. + :param subnet_id: The ID of the subnet in which to launch the instance. + :param security_group_ids: A list of security groups to associate with the notebook instance. + + :param lifecycle_config_name: The name of the lifecycle configuration to associate with the notebook + instance. + :param direct_internet_access: Whether to enable direct internet access for the notebook instance. + :param accelerator_types: The list of Elastic Inference (EI) accelerator types to associate with the + notebook instance. + :param default_code_repo: The URL of the Git repository that contains the default code for a notebook + instance. + :param additional_code_repos: A list of URLs for Git repositories that contain custom code for a + notebook instance. + :param root_access: Whether to give the notebook instance root access to the Amazon S3 bucket. + :param platform_id: The ID of the platform. + :param imds_config: The configuration for the instance metadata service. + :param config: Additional configuration options for the create call. + :param aws_conn_id: The AWS connection ID to use. + + :return: A dict containing the information about the create notebook instance call. + """ + if tags is not None: + tags = format_tags(tags) + + create_notebook_instance_kwargs = { + "NotebookInstanceName": instance_name, + "InstanceType": instance_type, + "RoleArn": role_arn, + "Tags": tags, + "VolumeSizeInGB": volume_size_in_gb, + "KmsKeyId": volume_kms_key_id, + "SubnetId": subnet_id, + "SecurityGroupIds": security_group_ids, + "LifecycleConfigName": lifecycle_config_name, + "DirectInternetAccess": direct_internet_access, + "AcceleratorTypes": accelerator_types, + "DefaultCodeRepository": default_code_repo, + "AdditionalCodeRepositories": additional_code_repos, + "RootAccess": root_access, + "PlatformIdentifier": platform_id, + } + + print(f"trimmed values are: {trim_none_values(create_notebook_instance_kwargs)}") + test = trim_none_values(create_notebook_instance_kwargs) + resp = self.conn.create_notebook_instance(**test, **kwargs) + return resp + + def start_instance(self, *, instance_name: str, **kwargs: Any) -> None: + """Start a notebook instance. + + :param instance_name: The name of the notebook instance. + :param wait_for_completion: Whether to wait for the start to complete. + :param kwargs: Additional arguments for the start notebook instance call. + :return: A dict containing the information about the start notebook instance call. + """ + resp = self.conn.start_notebook_instance(NotebookInstanceName=instance_name, **kwargs) + + return resp Review Comment: We try to avoid creating methods in hooks which add no value compare to the boto3 API. Here this method just wraps the boto3 API `start_notebook_instance` without doing anything else. You should remove this method and call the boto3 API directly from the operator using `hook.conn. start_notebook_instance(...)` ########## airflow/providers/amazon/aws/hooks/sagemaker_notebook.py: ########## @@ -0,0 +1,153 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import Any + +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook +from airflow.providers.amazon.aws.utils import trim_none_values +from airflow.providers.amazon.aws.utils.tags import format_tags + + +class SageMakerNotebookHook(AwsBaseHook): + """Interact with Amazon SageMaker to execute notebooks. + + Provide thick wrapper around + :external+boto3:py:class:`boto3.client('sagemaker') <SageMaker.Client>` + + Additional arguments (such as ``aws_conn_id``) may be specified and + are passed down to the underlying AwsBaseHook. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, *args, **kwargs) -> None: + kwargs["client_type"] = "sagemaker" + super().__init__(*args, **kwargs) + + def create_instance( Review Comment: Do you think it is worth exposing all these parameters? Usually we try to expose all mandatory parameters and parameters which we consider as important (very subjective). But here the list is quite long, I'd try to shorten it since you can still pass the parameters using `kwargs` ########## scripts/ci/docker-compose/empty/git_version: ########## Review Comment: I dont think this is needed? ########## tests/system/providers/amazon/aws/example_sagemaker_notebook.py: ########## @@ -0,0 +1,108 @@ +# 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_notebook 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 = f"{test_context[ENV_ID_KEY]}-test-notebook" + + # [START howto_operator_sagemaker_notebook_create] + instance = SageMakerCreateNotebookOperator( + task_id="create_instance", + instance_name=INSTANCE_NAME, + instance_type="ml.t3.medium", + role_arn=test_context[ROLE_ARN_KEY], + 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, + ) + Review Comment: ```suggestion ``` ########## tests/system/providers/amazon/aws/example_sagemaker_notebook.py: ########## @@ -0,0 +1,108 @@ +# 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_notebook 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 = f"{test_context[ENV_ID_KEY]}-test-notebook" + + # [START howto_operator_sagemaker_notebook_create] + instance = SageMakerCreateNotebookOperator( + task_id="create_instance", + instance_name=INSTANCE_NAME, + instance_type="ml.t3.medium", + role_arn=test_context[ROLE_ARN_KEY], + wait_for_completion=True, + ) + Review Comment: ```suggestion ``` ########## airflow/providers/amazon/aws/operators/sagemaker_notebook.py: ########## @@ -0,0 +1,287 @@ +# +# 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 functools import cached_property +from typing import TYPE_CHECKING, Sequence + +from airflow.models import BaseOperator +from airflow.providers.amazon.aws.hooks.sagemaker_notebook import SageMakerNotebookHook + +if TYPE_CHECKING: + pass + + +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 tags: A list of tags to associate with the notebook instance. + :param subnet_id: The ID of the subnet in which to launch the instance. + :param security_group_ids: A list of security groups to associate with the notebook instance. + :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 accelerator_types: The list of Elastic Inference (EI) accelerator types to associate with the + :param default_code_repo: The URL of the Git repository that contains the default code for a notebook + :param additional_code_repos: A list of URLs for Git repositories that contain custom code for a notebook + :param root_access: Whether to give the notebook instance root access to the Amazon S3 bucket. + :param platform_id: The ID of the platform. + :param imds_config: The configuration for the instance metadata service. + :param wait_for_completion: Whether or not to wait for the notebook to be InService before returning + :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", "config") Review Comment: As a convention, we try to include all parameters in `template_fields` ########## airflow/providers/amazon/aws/operators/sagemaker_notebook.py: ########## @@ -0,0 +1,287 @@ +# +# 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 functools import cached_property +from typing import TYPE_CHECKING, Sequence + +from airflow.models import BaseOperator +from airflow.providers.amazon.aws.hooks.sagemaker_notebook import SageMakerNotebookHook + +if TYPE_CHECKING: + pass + + +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 tags: A list of tags to associate with the notebook instance. + :param subnet_id: The ID of the subnet in which to launch the instance. + :param security_group_ids: A list of security groups to associate with the notebook instance. + :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 accelerator_types: The list of Elastic Inference (EI) accelerator types to associate with the + :param default_code_repo: The URL of the Git repository that contains the default code for a notebook + :param additional_code_repos: A list of URLs for Git repositories that contain custom code for a notebook + :param root_access: Whether to give the notebook instance root access to the Amazon S3 bucket. + :param platform_id: The ID of the platform. + :param imds_config: The configuration for the instance metadata service. + :param wait_for_completion: Whether or not to wait for the notebook to be InService before returning + :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", "config") + + ui_color = "#ff7300" + + def __init__( Review Comment: Same comment here on the parameters list length ########## airflow/providers/amazon/aws/operators/sagemaker_notebook.py: ########## @@ -0,0 +1,287 @@ +# +# 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 functools import cached_property +from typing import TYPE_CHECKING, Sequence + +from airflow.models import BaseOperator +from airflow.providers.amazon.aws.hooks.sagemaker_notebook import SageMakerNotebookHook + +if TYPE_CHECKING: + pass + + +class SageMakerCreateNotebookOperator(BaseOperator): + """ + Create a SageMaker notebook. Review Comment: An empty line is mandatory after the first line in docstring. ```suggestion Create a SageMaker notebook. ``` ########## airflow/providers/amazon/aws/hooks/sagemaker_notebook.py: ########## @@ -0,0 +1,153 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from typing import Any + +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook +from airflow.providers.amazon.aws.utils import trim_none_values +from airflow.providers.amazon.aws.utils.tags import format_tags + + +class SageMakerNotebookHook(AwsBaseHook): + """Interact with Amazon SageMaker to execute notebooks. + + Provide thick wrapper around + :external+boto3:py:class:`boto3.client('sagemaker') <SageMaker.Client>` + + Additional arguments (such as ``aws_conn_id``) may be specified and + are passed down to the underlying AwsBaseHook. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, *args, **kwargs) -> None: + kwargs["client_type"] = "sagemaker" + super().__init__(*args, **kwargs) + + def create_instance( + self, + instance_name: str, + instance_type: str, + role_arn: str, + volume_size_in_gb: int | None = None, + volume_kms_key_id: str | None = None, + tags: dict = {}, + subnet_id: str | None = None, + security_group_ids: list = [], + lifecycle_config_name: str | None = None, + direct_internet_access: str | None = None, + accelerator_types: list = [], + default_code_repo: list = [], + additional_code_repos: list = [], + root_access: str | None = None, + platform_id: str | None = None, + imds_config: dict | None = None, + **kwargs: Any, + ) -> dict: + """Create a SageMaker notebook instance. + + :param instance_name: The name of the notebook instance. + :param instance_type: The type of instance to create. For example, 'ml.t2.medium'. + :param image_uri: The Amazon EC2 Image URI for the SageMaker image to use. + :param role_arn: The Amazon Resource Name (ARN) of the IAM role to associate with the notebook + instance. + :param volume_size_in_gb: The size of the EBS volume to attach to the notebook instance. + :param volume_kms_key_id: The KMS key ID to use when creating the notebook instance. + :param tags: A list of tags to associate with the notebook instance. + :param subnet_id: The ID of the subnet in which to launch the instance. + :param security_group_ids: A list of security groups to associate with the notebook instance. + + :param lifecycle_config_name: The name of the lifecycle configuration to associate with the notebook + instance. + :param direct_internet_access: Whether to enable direct internet access for the notebook instance. + :param accelerator_types: The list of Elastic Inference (EI) accelerator types to associate with the + notebook instance. + :param default_code_repo: The URL of the Git repository that contains the default code for a notebook + instance. + :param additional_code_repos: A list of URLs for Git repositories that contain custom code for a + notebook instance. + :param root_access: Whether to give the notebook instance root access to the Amazon S3 bucket. + :param platform_id: The ID of the platform. + :param imds_config: The configuration for the instance metadata service. + :param config: Additional configuration options for the create call. + :param aws_conn_id: The AWS connection ID to use. + + :return: A dict containing the information about the create notebook instance call. + """ + if tags is not None: + tags = format_tags(tags) + + create_notebook_instance_kwargs = { + "NotebookInstanceName": instance_name, + "InstanceType": instance_type, + "RoleArn": role_arn, + "Tags": tags, + "VolumeSizeInGB": volume_size_in_gb, + "KmsKeyId": volume_kms_key_id, + "SubnetId": subnet_id, + "SecurityGroupIds": security_group_ids, + "LifecycleConfigName": lifecycle_config_name, + "DirectInternetAccess": direct_internet_access, + "AcceleratorTypes": accelerator_types, + "DefaultCodeRepository": default_code_repo, + "AdditionalCodeRepositories": additional_code_repos, + "RootAccess": root_access, + "PlatformIdentifier": platform_id, + } + + print(f"trimmed values are: {trim_none_values(create_notebook_instance_kwargs)}") + test = trim_none_values(create_notebook_instance_kwargs) + resp = self.conn.create_notebook_instance(**test, **kwargs) + return resp + + def start_instance(self, *, instance_name: str, **kwargs: Any) -> None: + """Start a notebook instance. + + :param instance_name: The name of the notebook instance. + :param wait_for_completion: Whether to wait for the start to complete. + :param kwargs: Additional arguments for the start notebook instance call. + :return: A dict containing the information about the start notebook instance call. + """ + resp = self.conn.start_notebook_instance(NotebookInstanceName=instance_name, **kwargs) + + return resp + + def stop_instance(self, *, instance_name: str, **kwargs: Any) -> None: + """Stop a notebook instance. + + :param instance_name: The name of the notebook instance. + :param wait_for_completion: Whether to wait for the stop to complete. + :param kwargs: Additional arguments for the stop notebook instance call. + :return: A dict containing the information about the stop notebook instance call. + """ + resp = self.conn.stop_notebook_instance(NotebookInstanceName=instance_name, **kwargs) + + return resp + + def delete_instance(self, *, instance_name: str, **kwargs: Any) -> None: + """Delete a notebook instance. + + :param instance_name: The name of the notebook instance. + :param kwargs: Additional arguments for the delete notebook instance call. + :return: A dict containing the information about the delete notebook instance call. + """ + resp = self.conn.delete_notebook_instance(NotebookInstanceName=instance_name, **kwargs) + + return resp Review Comment: Same here ########## tests/system/providers/amazon/aws/example_sagemaker_notebook.py: ########## @@ -0,0 +1,108 @@ +# 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_notebook 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 = f"{test_context[ENV_ID_KEY]}-test-notebook" + + # [START howto_operator_sagemaker_notebook_create] + instance = SageMakerCreateNotebookOperator( + task_id="create_instance", + instance_name=INSTANCE_NAME, + instance_type="ml.t3.medium", + role_arn=test_context[ROLE_ARN_KEY], + 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, + ) + Review Comment: ```suggestion ``` ########## tests/system/providers/amazon/aws/example_sagemaker_notebook.py: ########## @@ -0,0 +1,108 @@ +# 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_notebook 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 = f"{test_context[ENV_ID_KEY]}-test-notebook" + + # [START howto_operator_sagemaker_notebook_create] + instance = SageMakerCreateNotebookOperator( + task_id="create_instance", + instance_name=INSTANCE_NAME, + instance_type="ml.t3.medium", + role_arn=test_context[ROLE_ARN_KEY], + 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) + 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
