Merge branch 'master' into fix/missing-jumpstart-estimator-args

Author: evakravi

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## v2.171.0 (2023-07-06)
+
+### Features
+
+ * Add PipelineDefinitionConfig to pipelines to toggle custom job …
+
+### Bug Fixes and Other Changes
+
+ * Upgrade DJL deepspeed versions
+ * Remove unused dependency `protobuf3-to-dict`
+ * skip intelligent volume_size allocation based on instance type if it is a pipeline parameter
+
 ## v2.170.0 (2023-07-05)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.170.1.dev0
+2.171.1.dev0

src/sagemaker/clarify.py

@@ -267,6 +267,7 @@
                     },
                 },
                 SchemaOptional("seed"): int,
+                SchemaOptional("features_to_explain"): [Or(int, str)],
             },
             SchemaOptional("pre_training_bias"): {"methods": Or(str, [str])},
             SchemaOptional("post_training_bias"): {"methods": Or(str, [str])},
@@ -1308,6 +1309,7 @@ def __init__(
         num_clusters: Optional[int] = None,
         text_config: Optional[TextConfig] = None,
         image_config: Optional[ImageConfig] = None,
+        features_to_explain: Optional[List[Union[str, int]]] = None,
     ):
         """Initializes config for SHAP analysis.
 
@@ -1343,6 +1345,14 @@ def __init__(
                 text features. Default is None.
             image_config (:class:`~sagemaker.clarify.ImageConfig`): Config for handling image
                 features. Default is None.
+            features_to_explain: A list of names or indices of dataset features to compute SHAP
+                values for. If not provided, SHAP values are computed for all features by default.
+                Currently only supported for tabular datasets.
+
+        Raises:
+            ValueError: when ``agg_method`` is invalid, ``baseline`` and ``num_clusters`` are provided
+                together, or ``features_to_explain`` is specified when ``text_config`` or
+                ``image_config`` is provided
         """  # noqa E501  # pylint: disable=c0301
         if agg_method is not None and agg_method not in [
             "mean_abs",
@@ -1376,6 +1386,13 @@ def __init__(
                 )
         if image_config:
             _set(image_config.get_image_config(), "image_config", self.shap_config)
+        if features_to_explain is not None and (
+            text_config is not None or image_config is not None
+        ):
+            raise ValueError(
+                "`features_to_explain` is not supported for datasets containing text features or images."
+            )
+        _set(features_to_explain, "features_to_explain", self.shap_config)
 
     def get_explainability_config(self):
         """Returns a shap config dictionary."""

src/sagemaker/djl_inference/model.py

@@ -196,7 +196,7 @@ def _create_estimator(
     image_uri: str,
     role: str,
     sagemaker_session: Optional[Session],
-    volume_size: int = 30,
+    volume_size: int,
     vpc_config: Optional[
         Dict[
             str,
@@ -453,7 +453,9 @@ def partition(
         self,
         instance_type: str,
         s3_output_uri: str = None,
+        s3_output_prefix: str = "aot-partitioned-checkpoints",
         job_name: Optional[str] = None,
+        volume_size: int = 30,
         volume_kms_key: Optional[str] = None,
         output_kms_key: Optional[str] = None,
         use_spot_instances: bool = False,
@@ -469,8 +471,13 @@ def partition(
                     artifacts and output files). If not specified, results are
                     stored to a default bucket. If the bucket with the specific name
                     does not exist, it will be created.
+            s3_output_prefix (str): Name of the prefix where all the partitioned
+                    checkpoints to be uploaded. If not provided, the default value is
+                    aot-partitioned-checkpoints.
             job_name (str): Training job name. If not specified, a unique training job
                         name will be created.
+            volume_size (int): Size in GB of the storage volume to use for
+                storing input and output data during training (default: 30).
             volume_kms_key (str): Optional. KMS key ID for encrypting EBS
                 volume attached to the training instance (default: None).
             output_kms_key (str): Optional. KMS key ID for encrypting the
@@ -499,20 +506,19 @@ def partition(
             region_name = self.sagemaker_session.boto_session.region_name
             self.image_uri = self.serving_image_uri(region_name)
 
-        deploy_key_prefix = fw_utils.model_code_key_prefix(
-            self.key_prefix, self.name, self.image_uri
-        )
         if s3_output_uri is None:
+            deploy_key_prefix = fw_utils.model_code_key_prefix(
+                self.key_prefix, self.name, self.image_uri
+            )
+
             bucket, deploy_key_prefix = s3.determine_bucket_and_prefix(
                 bucket=self.bucket,
                 key_prefix=deploy_key_prefix,
                 sagemaker_session=self.sagemaker_session,
             )
             s3_output_uri = s3_path_join("s3://", bucket, deploy_key_prefix)
-        else:
-            s3_output_uri = s3_path_join(s3_output_uri, deploy_key_prefix)
 
-        self.save_mp_checkpoint_path = s3_path_join(s3_output_uri, "aot-partitioned-checkpoints")
+        self.save_mp_checkpoint_path = s3_path_join(s3_output_uri, s3_output_prefix)
 
         container_def = self._upload_model_to_s3(upload_as_tar=False)
         estimator = _create_estimator(
@@ -521,6 +527,7 @@ def partition(
             image_uri=self.image_uri,
             role=self.role,
             sagemaker_session=self.sagemaker_session,
+            volume_size=volume_size,
             vpc_config=self.vpc_config,
             volume_kms_key=volume_kms_key,
             output_kms_key=output_kms_key,
@@ -924,7 +931,9 @@ def partition(
         self,
         instance_type: str,
         s3_output_uri: str = None,
+        s3_output_prefix: str = "aot-partitioned-checkpoints",
         job_name: Optional[str] = None,
+        volume_size: int = 30,
         volume_kms_key: Optional[str] = None,
         output_kms_key: Optional[str] = None,
         use_spot_instances: bool = False,
@@ -940,8 +949,13 @@ def partition(
                     artifacts and output files). If not specified, results are
                     stored to a default bucket. If the bucket with the specific name
                     does not exist, it will be created.
+            s3_output_prefix (str): Name of the prefix where all the partitioned
+                    checkpoints to be uploaded. If not provided, the default value is
+                    aot-partitioned-checkpoints.
             job_name (str): Training job name. If not specified, a unique training job
                         name will be created.
+            volume_size (int): Size in GB of the storage volume to use for
+                storing input and output data during training (default: 30).
             volume_kms_key (str): Optional. KMS key ID for encrypting EBS
                 volume attached to the training instance (default: None).
             output_kms_key (str): Optional. KMS key ID for encrypting the
@@ -969,7 +983,9 @@ def partition(
         super(DeepSpeedModel, self).partition(
             instance_type,
             s3_output_uri,
-            job_name,
+            s3_output_prefix=s3_output_prefix,
+            job_name=job_name,
+            volume_size=volume_size,
             volume_kms_key=volume_kms_key,
             output_kms_key=output_kms_key,
             use_spot_instances=use_spot_instances,
@@ -1096,7 +1112,9 @@ def partition(
         self,
         instance_type: str,
         s3_output_uri: str = None,
+        s3_output_prefix: str = "aot-partitioned-checkpoints",
         job_name: Optional[str] = None,
+        volume_size: int = 30,
         volume_kms_key: Optional[str] = None,
         output_kms_key: Optional[str] = None,
         use_spot_instances: bool = False,
@@ -1112,8 +1130,13 @@ def partition(
                     artifacts and output files). If not specified, results are
                     stored to a default bucket. If the bucket with the specific name
                     does not exist, it will be created.
+            s3_output_prefix (str): Name of the prefix where all the partitioned
+                    checkpoints to be uploaded. If not provided, the default value is
+                    aot-partitioned-checkpoints.
             job_name (str): Training job name. If not specified, a unique training job
                         name will be created.
+            volume_size (int): Size in GB of the storage volume to use for
+                storing input and output data during training (default: 30).
             volume_kms_key (str): Optional. KMS key ID for encrypting EBS
                 volume attached to the training instance (default: None).
             output_kms_key (str): Optional. KMS key ID for encrypting the

tests/unit/test_clarify.py

@@ -716,6 +716,37 @@ def test_valid_shap_config(baseline):
     assert expected_config == shap_config.get_explainability_config()
 
 
+def test_shap_config_features_to_explain():
+    baseline = [1, 2, 3]
+    num_samples = 100
+    agg_method = "mean_sq"
+    use_logit = True
+    save_local_shap_values = True
+    seed = 123
+    features_to_explain = [0, 1]
+    shap_config = SHAPConfig(
+        baseline=baseline,
+        num_samples=num_samples,
+        agg_method=agg_method,
+        use_logit=use_logit,
+        save_local_shap_values=save_local_shap_values,
+        seed=seed,
+        features_to_explain=features_to_explain,
+    )
+    expected_config = {
+        "shap": {
+            "baseline": baseline,
+            "num_samples": num_samples,
+            "agg_method": agg_method,
+            "use_logit": use_logit,
+            "save_local_shap_values": save_local_shap_values,
+            "seed": seed,
+            "features_to_explain": features_to_explain,
+        }
+    }
+    assert expected_config == shap_config.get_explainability_config()
+
+
 def test_shap_config_no_baseline():
     num_samples = 100
     agg_method = "mean_sq"
@@ -852,6 +883,17 @@ def test_invalid_shap_config():
         "Baseline and num_clusters cannot be provided together. Please specify one of the two."
         in str(error.value)
     )
+    with pytest.raises(ValueError) as error:
+        SHAPConfig(
+            baseline=[[1, 2]],
+            num_samples=1,
+            text_config=TextConfig(granularity="token", language="english"),
+            features_to_explain=[0],
+        )
+    assert (
+        "`features_to_explain` is not supported for datasets containing text features or images."
+        in str(error.value)
+    )
 
 
 @pytest.fixture(scope="module")

tests/unit/test_djl_inference.py

@@ -598,7 +598,7 @@ def test_partition(
             IMAGE_URI, model_data_url="s3prefix", env=expected_env
         )
 
-        assert model.model_id == f"{s3_output_uri}s3prefix/aot-partitioned-checkpoints"
+        assert model.model_id == f"{s3_output_uri}aot-partitioned-checkpoints"
 
 
 @patch("sagemaker.djl_inference.model.fw_utils.model_code_key_prefix")
@@ -741,15 +741,15 @@ def without_user_input(sess):
             "s3://code-test-bucket/code-test-prefix/code-test-prefix-2",
             "s3://code-test-bucket/code-test-prefix/code-test-prefix-2/image_uri",
             "s3://code-test-bucket/code-test-prefix/code-test-prefix-2/image_uri",
-            "s3://test-bucket/test-prefix/test-prefix-2/code-test-prefix/code-test-prefix-2/image_uri",
-            "s3://test-bucket/test-prefix/test-prefix-2/code-test-prefix/code-test-prefix-2/image_uri",
+            "s3://test-bucket/test-prefix/test-prefix-2",
+            "s3://test-bucket/test-prefix/test-prefix-2",
         ),
         (
             None,
             f"s3://{DEFAULT_S3_BUCKET_NAME}/{DEFAULT_S3_OBJECT_KEY_PREFIX_NAME}/image_uri",
             f"s3://{DEFAULT_S3_BUCKET_NAME}/image_uri",
-            "s3://test-bucket/test-prefix/test-prefix-2/image_uri",
-            "s3://test-bucket/test-prefix/test-prefix-2/image_uri",
+            "s3://test-bucket/test-prefix/test-prefix-2",
+            "s3://test-bucket/test-prefix/test-prefix-2",
         ),
     ],
 )