feat: added endpoint_name to clarify.ModelConfig

Author: aws-byeldos

src/sagemaker/clarify.py

@@ -272,26 +272,30 @@ class ModelConfig:
 
     def __init__(
         self,
-        model_name,
-        instance_count,
-        instance_type,
-        accept_type=None,
-        content_type=None,
-        content_template=None,
-        custom_attributes=None,
-        accelerator_type=None,
-        endpoint_name_prefix=None,
-        target_model=None,
+        model_name: str = None,
+        instance_count: int = None,
+        instance_type: str = None,
+        accept_type: str = None,
+        content_type: str = None,
+        content_template: str = None,
+        custom_attributes: str = None,
+        accelerator_type: str = None,
+        endpoint_name_prefix: str = None,
+        target_model: str = None,
+        endpoint_name: str = None,
     ):
         r"""Initializes a configuration of a model and the endpoint to be created for it.
 
         Args:
             model_name (str): Model name (as created by
                 `CreateModel <https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateModel.html>`_.
+                Cannot be set when endpoint_name is set. Must be set with `instance_count`, `instance_type`
             instance_count (int): The number of instances of a new endpoint for model inference.
+                Cannot be set when endpoint_name is set. Must be set with `model_name`, `instance_type`
             instance_type (str): The type of
                 `EC2 instance <https://aws.amazon.com/ec2/instance-types/>`_
                 to use for model inference; for example, ``"ml.c5.xlarge"``.
+                Cannot be set when endpoint_name is set. Must be set with `instance_count`, `model_name`
             accept_type (str): The model output format to be used for getting inferences with the
                 shadow endpoint. Valid values are ``"text/csv"`` for CSV and
                 ``"application/jsonlines"``. Default is the same as ``content_type``.
@@ -321,17 +325,39 @@ def __init__(
             target_model (str): Sets the target model name when using a multi-model endpoint. For
                 more information about multi-model endpoints, see
                 https://docs.aws.amazon.com/sagemaker/latest/dg/multi-model-endpoints.html
+            endpoint_name (str): Sets the endpoint_name when re-uses an existing endpoint. Cannot be set
+                when `model_name`, `instance_count`, and `instance_type` set
 
         Raises:
-            ValueError: when the ``endpoint_name_prefix`` is invalid, ``accept_type`` is invalid,
-                 ``content_type`` is invalid, or ``content_template`` has no placeholder "features"
+            ValueError: when the
+                - ``endpoint_name_prefix`` is invalid,
+                - ``accept_type`` is invalid,
+                - ``content_type`` is invalid,
+                - ``content_template`` has no placeholder "features"
+                - both [``endpoint_name``] AND [``model_name``, ``instance_count``, ``instance_type``] are set
+                - both [``endpoint_name``] AND [``endpoint_name_prefix``] are set
         """
-        self.predictor_config = {
-            "model_name": model_name,
-            "instance_type": instance_type,
-            "initial_instance_count": instance_count,
-        }
-        if endpoint_name_prefix is not None:
+
+        # validation
+        _model_endpoint_config_rule = (
+            all([model_name, instance_count, instance_type]),
+            all([endpoint_name]),
+        )
+        assert any(_model_endpoint_config_rule) and not all(_model_endpoint_config_rule)
+        if endpoint_name:
+            assert not endpoint_name_prefix
+
+        # main init logic
+        self.predictor_config = (
+            {
+                "model_name": model_name,
+                "instance_type": instance_type,
+                "initial_instance_count": instance_count,
+            }
+            if not endpoint_name
+            else {"endpoint_name": endpoint_name}
+        )
+        if endpoint_name_prefix:
             if re.search("^[a-zA-Z0-9](-*[a-zA-Z0-9])", endpoint_name_prefix) is None:
                 raise ValueError(
                     "Invalid endpoint_name_prefix."

tests/unit/test_clarify.py

@@ -735,6 +735,41 @@ def pdp_config():
     return PDPConfig(features=["F1", "F2"], grid_resolution=20)
 
 
+def test_model_config_validations():
+    new_model_endpoint_definition = {
+        "model_name": "xgboost-model",
+        "instance_type": "ml.c5.xlarge",
+        "instance_count": 1,
+    }
+    existing_endpoint_definition = {"endpoint_name": "existing_endpoint"}
+
+    with pytest.raises(AssertionError):
+        # should be one of them
+        ModelConfig(
+            **new_model_endpoint_definition,
+            **existing_endpoint_definition,
+        )
+
+    with pytest.raises(AssertionError):
+        # should be one of them
+        ModelConfig(
+            endpoint_name_prefix="prefix",
+            **existing_endpoint_definition,
+        )
+
+    # success path for new model
+    assert ModelConfig(**new_model_endpoint_definition).predictor_config == {
+        "initial_instance_count": 1,
+        "instance_type": "ml.c5.xlarge",
+        "model_name": "xgboost-model",
+    }
+
+    # success path for existing endpoint
+    assert (
+        ModelConfig(**existing_endpoint_definition).predictor_config == existing_endpoint_definition
+    )
+
+
 @patch("sagemaker.utils.name_from_base", return_value=JOB_NAME)
 def test_pre_training_bias(
     name_from_base,
@@ -1396,6 +1431,47 @@ def test_analysis_config_generator_for_bias_explainability(
     assert actual == expected
 
 
+def test_analysis_config_generator_for_bias_explainability_with_existing_endpoint(
+    data_config, data_bias_config
+):
+    model_config = ModelConfig(endpoint_name="existing_endpoint_name")
+    model_predicted_label_config = ModelPredictedLabelConfig(
+        probability="pr",
+        label_headers=["success"],
+    )
+    actual = _AnalysisConfigGenerator.bias_and_explainability(
+        data_config,
+        model_config,
+        model_predicted_label_config,
+        [SHAPConfig(), PDPConfig()],
+        data_bias_config,
+        pre_training_methods="all",
+        post_training_methods="all",
+    )
+    expected = {
+        "dataset_type": "text/csv",
+        "facet": [{"name_or_index": "F1"}],
+        "group_variable": "F2",
+        "headers": ["Label", "F1", "F2", "F3", "F4"],
+        "joinsource_name_or_index": "F4",
+        "label": "Label",
+        "label_values_or_threshold": [1],
+        "methods": {
+            "pdp": {"grid_resolution": 15, "top_k_features": 10},
+            "post_training_bias": {"methods": "all"},
+            "pre_training_bias": {"methods": "all"},
+            "report": {"name": "report", "title": "Analysis Report"},
+            "shap": {"save_local_shap_values": True, "use_logit": False},
+        },
+        "predictor": {
+            "label_headers": ["success"],
+            "endpoint_name": "existing_endpoint_name",
+            "probability": "pr",
+        },
+    }
+    assert actual == expected
+
+
 def test_analysis_config_generator_for_bias_pre_training(data_config, data_bias_config):
     actual = _AnalysisConfigGenerator.bias_pre_training(
         data_config, data_bias_config, methods="all"