Merge branch 'master' into 3826-bugfix-lambda-step

Author: dipanjank

.gitignore

@@ -30,6 +30,7 @@ env/
 .vscode/
 **/tmp
 .python-version
+*.html
 **/_repack_script_launcher.sh
 tests/data/**/_repack_model.py
 tests/data/experiment/sagemaker-dev-1.0.tar.gz

CHANGELOG.md

@@ -1,5 +1,30 @@
 # Changelog
 
+## v2.152.0 (2023-05-04)
+
+### Features
+
+ * add support for lineage visualization using pyvis
+ * Expose Experiment class publicly
+ * PyTorch 1.13 release
+
+### Bug Fixes and Other Changes
+
+ * Change data_type argument to dtype to keep consistent with D…
+ * Skip edge test
+ * make RemoteExecutor context manager non-blocking on pending futures
+ * Add inferentia2 DLC images for djl framework
+ * Fix typo in using_pytorch.rst
+ * Unable to attach estimator to training job when KeepAlivePeriodInSeconds specified
+ * update LMI container image
+ * Update Clarify SHAPConfig baseline to allow JSON structures
+
+### Documentation Changes
+
+ * Fix broken link in DJL SageMaker docs
+ * currency update for the SageMaker data parallelism lib
+ * SM model parallel library v1.15.0 release note
+
 ## v2.151.0 (2023-04-27)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.151.1.dev0
+2.152.1.dev0

doc/api/training/sdp_versions/latest.rst

@@ -26,7 +26,7 @@ depending on the version of the library you use.
    <https://docs.aws.amazon.com/sagemaker/latest/dg/data-parallel-use-api.html#data-parallel-use-python-skd-api>`_
    for more information.
 
-For versions between 1.4.0 and 1.7.0 (Latest)
+For versions between 1.4.0 and 1.8.0 (Latest)
 =============================================
 
 .. toctree::

doc/api/training/smd_data_parallel_release_notes/smd_data_parallel_change_log.rst

@@ -5,39 +5,64 @@ Release Notes
 #############
 
 New features, bug fixes, and improvements are regularly made to the SageMaker
-distributed data parallel library.
+data parallelism library.
 
-SageMaker Distributed Data Parallel 1.7.0 Release Notes
+SageMaker Distributed Data Parallel 1.8.0 Release Notes
 =======================================================
 
-*Date: Feb. 10. 2023*
+*Date: Apr. 17. 2023*
 
 **Currency Updates**
 
-* Added support for PyTorch 1.13.1.
+* Added support for PyTorch 2.0.0.
 
 **Migration to AWS Deep Learning Containers**
 
 This version passed benchmark testing and is migrated to the following AWS Deep Learning Containers (DLC):
 
-- PyTorch 1.13.1 DLC
+- PyTorch 2.0.0 DLC
 
   .. code::
 
-    763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.13.1-gpu-py39-cu117-ubuntu20.04-sagemaker
+    763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker
 
 Binary file of this version of the library for custom container users:
 
   .. code::
 
-    https://smdataparallel.s3.amazonaws.com/binary/pytorch/1.13.1/cu117/2023-01-09/smdistributed_dataparallel-1.7.0-cp39-cp39-linux_x86_64.whl
+    https://smdataparallel.s3.amazonaws.com/binary/pytorch/2.0.0/cu118/2023-03-20/smdistributed_dataparallel-1.8.0-cp310-cp310-linux_x86_64.whl
 
 
 ----
 
 Release History
 ===============
 
+SageMaker Distributed Data Parallel 1.7.0 Release Notes
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+*Date: Feb. 10. 2023*
+
+**Currency Updates**
+
+* Added support for PyTorch 1.13.1.
+
+**Migration to AWS Deep Learning Containers**
+
+This version passed benchmark testing and is migrated to the following AWS Deep Learning Containers (DLC):
+
+- PyTorch 1.13.1 DLC
+
+  .. code::
+
+    763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.13.1-gpu-py39-cu117-ubuntu20.04-sagemaker
+
+Binary file of this version of the library for custom container users:
+
+  .. code::
+
+    https://smdataparallel.s3.amazonaws.com/binary/pytorch/1.13.1/cu117/2023-01-09/smdistributed_dataparallel-1.7.0-cp39-cp39-linux_x86_64.whl
+
 SageMaker Distributed Data Parallel 1.6.0 Release Notes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

doc/experiments/sagemaker.experiments.rst

@@ -11,6 +11,15 @@ Run
 
 .. automethod:: sagemaker.experiments.list_runs
 
+Experiment
+-------------
+
+.. autoclass:: sagemaker.experiments.Experiment
+    :members:
+
+Other
+-------------
+
 .. autoclass:: sagemaker.experiments.SortByType
     :members:
     :undoc-members:

doc/frameworks/djl/using_djl.rst

@@ -31,7 +31,7 @@ You can either deploy your model using DeepSpeed or HuggingFace Accelerate, or l
     djl_model = DJLModel(
         "s3://my_bucket/my_saved_model_artifacts/", # This can also be a HuggingFace Hub model id
         "my_sagemaker_role",
-        data_type="fp16",
+        dtype="fp16",
         task="text-generation",
         number_of_partitions=2 # number of gpus to partition the model across
     )
@@ -48,7 +48,7 @@ If you want to use a specific backend, then you can create an instance of the co
     deepspeed_model = DeepSpeedModel(
         "s3://my_bucket/my_saved_model_artifacts/", # This can also be a HuggingFace Hub model id
         "my_sagemaker_role",
-        data_type="bf16",
+        dtype="bf16",
         task="text-generation",
         tensor_parallel_degree=2, # number of gpus to partition the model across using tensor parallelism
     )
@@ -58,7 +58,7 @@ If you want to use a specific backend, then you can create an instance of the co
     hf_accelerate_model = HuggingFaceAccelerateModel(
         "s3://my_bucket/my_saved_model_artifacts/", # This can also be a HuggingFace Hub model id
         "my_sagemaker_role",
-        data_type="fp16",
+        dtype="fp16",
         task="text-generation",
         number_of_partitions=2, # number of gpus to partition the model across
     )
@@ -109,7 +109,7 @@ For example, you can deploy the EleutherAI gpt-j-6B model like this:
     model = DJLModel(
         "EleutherAI/gpt-j-6B",
         "my_sagemaker_role",
-        data_type="fp16",
+        dtype="fp16",
         number_of_partitions=2
     )
 
@@ -142,7 +142,7 @@ You would then pass "s3://my_bucket/gpt-j-6B" as ``model_id`` to the ``DJLModel`
     model = DJLModel(
         "s3://my_bucket/gpt-j-6B",
         "my_sagemaker_role",
-        data_type="fp16",
+        dtype="fp16",
         number_of_partitions=2
     )
 
@@ -213,7 +213,7 @@ For more information about DJL Serving, see the `DJL Serving documentation. <htt
 SageMaker DJL Classes
 ***********************
 
-For information about the different DJL Serving related classes in the SageMaker Python SDK, see https://sagemaker.readthedocs.io/en/stable/sagemaker.djl_inference.html.
+For information about the different DJL Serving related classes in the SageMaker Python SDK, see https://sagemaker.readthedocs.io/en/stable/frameworks/djl/sagemaker.djl_inference.html.
 
 ********************************
 SageMaker DJL Serving Containers

requirements/extras/test_requirements.txt

@@ -19,6 +19,7 @@ fabric==2.6.0
 requests==2.27.1
 sagemaker-experiments==0.1.35
 Jinja2==3.0.3
+pyvis==0.2.1
 pandas>=1.3.5,<1.5
 scikit-learn==1.0.2
 cloudpickle==2.2.1

src/sagemaker/djl_inference/model.py

@@ -233,7 +233,7 @@ def __init__(
         role: str,
         djl_version: Optional[str] = None,
         task: Optional[str] = None,
-        data_type: str = "fp32",
+        dtype: str = "fp32",
         number_of_partitions: Optional[int] = None,
         min_workers: Optional[int] = None,
         max_workers: Optional[int] = None,
@@ -264,7 +264,7 @@ def __init__(
             task (str): The HuggingFace/NLP task you want to launch this model for. Defaults to
                 None.
                 If not provided, the task will be inferred from the model architecture by DJL.
-            data_type (str): The data type to use for loading your model. Accepted values are
+            dtype (str): The data type to use for loading your model. Accepted values are
                 "fp32", "fp16", "bf16", "int8". Defaults to "fp32".
             number_of_partitions (int): The number of GPUs to partition the model across. The
                 partitioning strategy is determined by the selected backend. If DeepSpeed is
@@ -322,13 +322,20 @@ def __init__(
                 "You only need to set model_id and ensure it points to uncompressed model "
                 "artifacts in s3, or a valid HuggingFace Hub model_id."
             )
+        data_type = kwargs.pop("data_type", None)
+        if data_type:
+            logger.warning(
+                "data_type is being deprecated in favor of dtype. Please migrate use of data_type"
+                " to dtype. Support for data_type will be removed in a future release"
+            )
+            dtype = dtype or data_type
         super(DJLModel, self).__init__(
             None, image_uri, role, entry_point, predictor_cls=predictor_cls, **kwargs
         )
         self.model_id = model_id
         self.djl_version = djl_version
         self.task = task
-        self.data_type = data_type
+        self.dtype = dtype
         self.number_of_partitions = number_of_partitions
         self.min_workers = min_workers
         self.max_workers = max_workers
@@ -372,7 +379,7 @@ def transformer(self, **_):
             "DJLModels do not currently support Batch Transform inference jobs"
         )
 
-    def right_size(self, checkpoint_data_type: str):
+    def right_size(self, **_):
         """Not implemented.
 
         DJLModels do not support SageMaker Inference Recommendation Jobs.
@@ -573,8 +580,8 @@ def generate_serving_properties(self, serving_properties=None) -> Dict[str, str]
             serving_properties["option.entryPoint"] = self.entry_point
         if self.task:
             serving_properties["option.task"] = self.task
-        if self.data_type:
-            serving_properties["option.dtype"] = self.data_type
+        if self.dtype:
+            serving_properties["option.dtype"] = self.dtype
         if self.min_workers:
             serving_properties["minWorkers"] = self.min_workers
         if self.max_workers:
@@ -779,7 +786,7 @@ def __init__(
                 None.
             load_in_8bit (bool): Whether to load the model in int8 precision using bits and bytes
                 quantization. This is only supported for select model architectures.
-                Defaults to False. If ``data_type`` is int8, then this is set to True.
+                Defaults to False. If ``dtype`` is int8, then this is set to True.
             low_cpu_mem_usage (bool): Whether to limit CPU memory usage to 1x model size during
                 model loading. This is an experimental feature in HuggingFace. This is useful when
                 loading multiple instances of your model in parallel. Defaults to False.
@@ -832,19 +839,19 @@ def generate_serving_properties(self, serving_properties=None) -> Dict[str, str]
         if self.device_map:
             serving_properties["option.device_map"] = self.device_map
         if self.load_in_8bit:
-            if self.data_type != "int8":
-                raise ValueError("Set data_type='int8' to use load_in_8bit")
+            if self.dtype != "int8":
+                raise ValueError("Set dtype='int8' to use load_in_8bit")
             serving_properties["option.load_in_8bit"] = self.load_in_8bit
-        if self.data_type == "int8":
+        if self.dtype == "int8":
             serving_properties["option.load_in_8bit"] = True
         if self.low_cpu_mem_usage:
             serving_properties["option.low_cpu_mem_usage"] = self.low_cpu_mem_usage
         # This is a workaround due to a bug in our built in handler for huggingface
         # TODO: This needs to be fixed when new dlc is published
         if (
             serving_properties["option.entryPoint"] == "djl_python.huggingface"
-            and self.data_type
-            and self.data_type != "auto"
+            and self.dtype
+            and self.dtype != "auto"
         ):
             serving_properties["option.dtype"] = "auto"
             serving_properties.pop("option.load_in_8bit", None)

src/sagemaker/experiments/__init__.py

@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 
 from sagemaker.experiments.run import Run  # noqa: F401
+from sagemaker.experiments.experiment import Experiment  # noqa: F401
 from sagemaker.experiments.run import load_run  # noqa: F401
 from sagemaker.experiments.run import list_runs  # noqa: F401
 from sagemaker.experiments.run import SortOrderType  # noqa: F401

src/sagemaker/experiments/experiment.py

@@ -22,11 +22,11 @@
 from sagemaker.experiments.trial_component import _TrialComponent
 
 
-class _Experiment(_base_types.Record):
+class Experiment(_base_types.Record):
     """An Amazon SageMaker experiment, which is a collection of related trials.
 
-    New experiments are created by calling `experiments.experiment._Experiment.create`.
-    Existing experiments can be reloaded by calling `experiments.experiment._Experiment.load`.
+    New experiments are created by calling `experiments.experiment.Experiment.create`.
+    Existing experiments can be reloaded by calling `experiments.experiment.Experiment.load`.
 
     Attributes:
         experiment_name (str): The name of the experiment. The name must be unique
@@ -73,7 +73,7 @@ def delete(self):
 
     @classmethod
     def load(cls, experiment_name, sagemaker_session=None):
-        """Load an existing experiment and return an `_Experiment` object representing it.
+        """Load an existing experiment and return an `Experiment` object representing it.
 
         Args:
             experiment_name: (str): Name of the experiment
@@ -83,7 +83,7 @@ def load(cls, experiment_name, sagemaker_session=None):
                 default AWS configuration chain.
 
         Returns:
-            experiments.experiment._Experiment: A SageMaker `_Experiment` object
+            experiments.experiment.Experiment: A SageMaker `Experiment` object
         """
         return cls._construct(
             cls._boto_load_method,
@@ -100,7 +100,7 @@ def create(
         tags=None,
         sagemaker_session=None,
     ):
-        """Create a new experiment in SageMaker and return an `_Experiment` object.
+        """Create a new experiment in SageMaker and return an `Experiment` object.
 
         Args:
             experiment_name: (str): Name of the experiment. Must be unique. Required.
@@ -115,7 +115,7 @@ def create(
                 (default: None).
 
         Returns:
-            experiments.experiment._Experiment: A SageMaker `_Experiment` object
+            experiments.experiment.Experiment: A SageMaker `Experiment` object
         """
         return cls._construct(
             cls._boto_create_method,
@@ -154,10 +154,10 @@ def _load_or_create(
                 exist and a new experiment has to be created.
 
         Returns:
-            experiments.experiment._Experiment: A SageMaker `_Experiment` object
+            experiments.experiment.Experiment: A SageMaker `Experiment` object
         """
         try:
-            experiment = _Experiment.create(
+            experiment = Experiment.create(
                 experiment_name=experiment_name,
                 display_name=display_name,
                 description=description,
@@ -170,7 +170,7 @@ def _load_or_create(
             if not (error_code == "ValidationException" and "already exists" in error_message):
                 raise ce
             # already exists
-            experiment = _Experiment.load(experiment_name, sagemaker_session)
+            experiment = Experiment.load(experiment_name, sagemaker_session)
         return experiment
 
     def list_trials(self, created_before=None, created_after=None, sort_by=None, sort_order=None):

src/sagemaker/experiments/run.py

@@ -32,7 +32,7 @@
 )
 from sagemaker.experiments._environment import _RunEnvironment
 from sagemaker.experiments._run_context import _RunContext
-from sagemaker.experiments.experiment import _Experiment
+from sagemaker.experiments.experiment import Experiment
 from sagemaker.experiments._metrics import _MetricsManager
 from sagemaker.experiments.trial import _Trial
 from sagemaker.experiments.trial_component import _TrialComponent
@@ -166,7 +166,7 @@ def __init__(
         )
         self.run_group_name = Run._generate_trial_name(self.experiment_name)
 
-        self._experiment = _Experiment._load_or_create(
+        self._experiment = Experiment._load_or_create(
             experiment_name=self.experiment_name,
             display_name=experiment_display_name,
             tags=tags,