Add example notebook (#1528)

* Add testing notebook

* format

* use smaller data

* remove large dataset

* update

* pylint

* flake8

* ignore docstyle in directories with test

* format

* format

Author: benieric

.pydocstylerc

@@ -2,3 +2,4 @@
 inherit = false
 ignore = D104,D107,D202,D203,D213,D214,D400,D401,D404,D406,D407,D411,D413,D414,D415,D417
 match = (?!record_pb2).*\.py
+match-dir = (?!.*test).*

src/sagemaker/modules/image_spec.py

@@ -16,17 +16,28 @@
 import re
 from enum import Enum
 from typing import Optional
+from packaging.version import Version
 
 from sagemaker import utils
-from sagemaker.image_uris import _validate_version_and_set_if_needed, _version_for_config, \
-    _config_for_framework_and_scope, _validate_py_version_and_set_if_needed, _registry_from_region, ECR_URI_TEMPLATE, \
-    _get_latest_versions, _validate_instance_deprecation, _get_image_tag, _validate_arg
-from packaging.version import Version
+from sagemaker.image_uris import (
+    _validate_version_and_set_if_needed,
+    _version_for_config,
+    _config_for_framework_and_scope,
+    _validate_py_version_and_set_if_needed,
+    _registry_from_region,
+    ECR_URI_TEMPLATE,
+    _get_latest_versions,
+    _validate_instance_deprecation,
+    _get_image_tag,
+    _validate_arg,
+)
 
 DEFAULT_TOLERATE_MODEL = False
 
 
 class Framework(Enum):
+    """Framework enum class."""
+
     HUGGING_FACE = "huggingface"
     HUGGING_FACE_NEURON = "huggingface-neuron"
     HUGGING_FACE_NEURON_X = "huggingface-neuronx"
@@ -46,12 +57,16 @@ class Framework(Enum):
 
 
 class ImageScope(Enum):
+    """ImageScope enum class."""
+
     TRAINING = "training"
     INFERENCE = "inference"
     INFERENCE_GRAVITON = "inference-graviton"
 
 
 class Processor(Enum):
+    """Processor enum class."""
+
     INF = "inf"
     NEURON = "neuron"
     GPU = "gpu"
@@ -60,22 +75,53 @@ class Processor(Enum):
 
 
 class ImageSpec:
-    """ImageSpec class to get image URI for a specific framework version."""
-
-    def __init__(self,
-                 framework: Framework,
-                 processor: Optional[Processor] = Processor.CPU,
-                 region: Optional[str] = "us-west-2",
-                 version=None,
-                 py_version=None,
-                 instance_type=None,
-                 accelerator_type=None,
-                 image_scope: ImageScope = ImageScope.TRAINING,
-                 container_version=None,
-                 distribution=None,
-                 base_framework_version=None,
-                 sdk_version=None,
-                 inference_tool=None):
+    """ImageSpec class to get image URI for a specific framework version.
+
+    Attributes:
+        framework (Framework): The name of the framework or algorithm.
+        processor (Processor): The name of the processor (CPU, GPU, etc.).
+        region (str): The AWS region.
+        version (str): The framework or algorithm version. This is required if there is
+            more than one supported version for the given framework or algorithm.
+        py_version (str): The Python version. This is required if there is
+            more than one supported Python version for the given framework version.
+        instance_type (str): The SageMaker instance type. For supported types, see
+            https://aws.amazon.com/sagemaker/pricing. This is required if
+            there are different images for different processor types.
+        accelerator_type (str): Elastic Inference accelerator type. For more, see
+            https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.
+        image_scope (str): The image type, i.e. what it is used for.
+            Valid values: "training", "inference", "inference_graviton", "eia".
+            If ``accelerator_type`` is set, ``image_scope`` is ignored.
+        container_version (str): the version of docker image.
+            Ideally the value of parameter should be created inside the framework.
+            For custom use, see the list of supported container versions:
+            https://github.com/aws/deep-learning-containers/blob/master/available_images.md
+            (default: None).
+        distribution (dict): A dictionary with information on how to run distributed training
+        sdk_version (str): the version of python-sdk that will be used in the image retrieval.
+            (default: None).
+        inference_tool (str): the tool that will be used to aid in the inference.
+            Valid values: "neuron, neuronx, None"
+            (default: None).
+    """
+
+    def __init__(
+        self,
+        framework: Framework,
+        processor: Optional[Processor] = Processor.CPU,
+        region: Optional[str] = "us-west-2",
+        version=None,
+        py_version=None,
+        instance_type=None,
+        accelerator_type=None,
+        image_scope: ImageScope = ImageScope.TRAINING,
+        container_version=None,
+        distribution=None,
+        base_framework_version=None,
+        sdk_version=None,
+        inference_tool=None,
+    ):
         self.framework = framework
         self.processor = processor
         self.version = version
@@ -91,44 +137,14 @@ def __init__(self,
         self.inference_tool = inference_tool
 
     def update_image_spec(self, **kwargs):
+        """Update the ImageSpec object with the given arguments."""
         for key, value in kwargs.items():
             if hasattr(self, key):
                 setattr(self, key, value)
 
     def retrieve(self) -> str:
         """Retrieves the ECR URI for the Docker image matching the given arguments.
 
-        Ideally this function should not be called directly, rather it should be called from the
-        fit() function inside framework estimator.
-
-        Args:
-            framework (Framework): The name of the framework or algorithm.
-            processor (Processor): The name of the processor (CPU, GPU, etc.).
-            region (str): The AWS region.
-            version (str): The framework or algorithm version. This is required if there is
-                more than one supported version for the given framework or algorithm.
-            py_version (str): The Python version. This is required if there is
-                more than one supported Python version for the given framework version.
-            instance_type (str): The SageMaker instance type. For supported types, see
-                https://aws.amazon.com/sagemaker/pricing. This is required if
-                there are different images for different processor types.
-            accelerator_type (str): Elastic Inference accelerator type. For more, see
-                https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.
-            image_scope (str): The image type, i.e. what it is used for.
-                Valid values: "training", "inference", "inference_graviton", "eia".
-                If ``accelerator_type`` is set, ``image_scope`` is ignored.
-            container_version (str): the version of docker image.
-                Ideally the value of parameter should be created inside the framework.
-                For custom use, see the list of supported container versions:
-                https://github.com/aws/deep-learning-containers/blob/master/available_images.md
-                (default: None).
-            distribution (dict): A dictionary with information on how to run distributed training
-            sdk_version (str): the version of python-sdk that will be used in the image retrieval.
-                (default: None).
-            inference_tool (str): the tool that will be used to aid in the inference.
-                Valid values: "neuron, neuronx, None"
-                (default: None).
-
         Returns:
             str: The ECR URI for the corresponding SageMaker Docker image.
 
@@ -140,13 +156,14 @@ def retrieve(self) -> str:
                 known security vulnerabilities.
             DeprecatedJumpStartModelError: If the version of the model is deprecated.
         """
-        config = _config_for_framework_and_scope(self.framework.value,
-                                                 self.image_scope.value,
-                                                 self.accelerator_type)
-
+        config = _config_for_framework_and_scope(
+            self.framework.value, self.image_scope.value, self.accelerator_type
+        )
         original_version = self.version
         try:
-            version = _validate_version_and_set_if_needed(self.version, config, self.framework.value)
+            version = _validate_version_and_set_if_needed(
+                self.version, config, self.framework.value
+            )
         except ValueError:
             version = None
         if not version:
@@ -159,12 +176,14 @@ def retrieve(self) -> str:
                 full_base_framework_version = version_config["version_aliases"].get(
                     self.base_framework_version, self.base_framework_version
                 )
-            _validate_arg(full_base_framework_version, list(version_config.keys()), "base framework")
+            _validate_arg(
+                full_base_framework_version, list(version_config.keys()), "base framework"
+            )
             version_config = version_config.get(full_base_framework_version)
 
-        self.py_version = _validate_py_version_and_set_if_needed(self.py_version,
-                                                            version_config,
-                                                            self.framework.value)
+        self.py_version = _validate_py_version_and_set_if_needed(
+            self.py_version, version_config, self.framework.value
+        )
         version_config = version_config.get(self.py_version) or version_config
 
         registry = _registry_from_region(self.region, version_config["registries"])
@@ -206,16 +225,18 @@ def retrieve(self) -> str:
                 if config.get("version_aliases").get(original_version):
                     _version = config.get("version_aliases")[original_version]
                 if (
-                        config.get("versions", {})
-                                .get(_version, {})
-                                .get("version_aliases", {})
-                                .get(self.base_framework_version, {})
+                    config.get("versions", {})
+                    .get(_version, {})
+                    .get("version_aliases", {})
+                    .get(self.base_framework_version, {})
                 ):
                     _base_framework_version = config.get("versions")[_version]["version_aliases"][
                         self.base_framework_version
                     ]
                     pt_or_tf_version = (
-                        re.compile("^(pytorch|tensorflow)(.*)$").match(_base_framework_version).group(2)
+                        re.compile("^(pytorch|tensorflow)(.*)$")
+                        .match(_base_framework_version)
+                        .group(2)
                     )
 
             tag_prefix = f"{pt_or_tf_version}-transformers{_version}"
@@ -224,29 +245,28 @@ def retrieve(self) -> str:
 
         if repo == f"{self.framework.value}-inference-graviton":
             self.container_version = f"{self.container_version}-sagemaker"
-        _validate_instance_deprecation(self.framework,
-                                       self.instance_type,
-                                       version)
+        _validate_instance_deprecation(self.framework, self.instance_type, version)
 
         tag = _get_image_tag(
             self.container_version,
             self.distribution,
             self.image_scope.value,
-            self.framework,
+            self.framework.value,
             self.inference_tool,
             self.instance_type,
             self.processor.value,
             self.py_version,
             tag_prefix,
-            version)
+            version,
+        )
 
         if tag:
             repo += ":{}".format(tag)
 
         return ECR_URI_TEMPLATE.format(registry=registry, hostname=hostname, repository=repo)
 
-    def _fetch_latest_version_from_config(self,
-                                          framework_config: dict) -> str:
+    def _fetch_latest_version_from_config(self, framework_config: dict) -> str:
+        """Fetches the latest version from the framework config."""
         if self.image_scope.value in framework_config:
             if image_scope_config := framework_config[self.image_scope.value]:
                 if version_aliases := image_scope_config["version_aliases"]: