fix: fix HuggingFace GEN2 model deployment arguments (#1404)

Author: bhaoz

src/sagemaker/huggingface/model.py

@@ -331,6 +331,9 @@ def deploy(
             inference_recommendation_id=inference_recommendation_id,
             explainer_config=explainer_config,
             endpoint_logging=kwargs.get("endpoint_logging", False),
+            endpoint_type=kwargs.get("endpoint_type", None),
+            resources=kwargs.get("resources", None),
+            managed_instance_scaling=kwargs.get("managed_instance_scaling", None),
         )
 
     def register(

tests/integ/test_huggingface.py

@@ -21,6 +21,8 @@
 from sagemaker.utils import unique_name_from_base
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
+from sagemaker.enums import EndpointType
+from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
 
 ROLE = "SageMakerRole"
 
@@ -172,3 +174,57 @@ def test_huggingface_inference(
         }
         output = predictor.predict(data)
         assert "score" in output[0]
+
+
+@pytest.mark.skip(
+    reason="re-enable when above GEN1 endpoint hugging face inference test enabled",
+)
+def test_huggingface_inference_gen2_endpoint(
+    sagemaker_session,
+    gpu_pytorch_instance_type,
+    huggingface_inference_latest_version,
+    huggingface_inference_pytorch_latest_version,
+    huggingface_pytorch_latest_inference_py_version,
+):
+    env = {
+        "HF_MODEL_ID": "philschmid/tiny-distilbert-classification",
+        "HF_TASK": "text-classification",
+    }
+    endpoint_name = unique_name_from_base("test-hf-inference")
+
+    model = HuggingFaceModel(
+        sagemaker_session=sagemaker_session,
+        role="SageMakerRole",
+        env=env,
+        py_version=huggingface_pytorch_latest_inference_py_version,
+        transformers_version=huggingface_inference_latest_version,
+        pytorch_version=huggingface_inference_pytorch_latest_version,
+    )
+    predictor = model.deploy(
+        instance_type=gpu_pytorch_instance_type,
+        initial_instance_count=1,
+        endpoint_name=endpoint_name,
+        endpoint_type=EndpointType.GEN2,
+        resources=ResourceRequirements(
+            requests={
+                "num_accelerators": 1,  # NumberOfCpuCoresRequired
+                "memory": 8192,  # MinMemoryRequiredInMb (required)
+                "copies": 1,
+            },
+            limits={},
+        ),
+    )
+
+    data = {
+        "inputs": "Camera - You are awarded a SiPix Digital Camera!"
+        "call 09061221066 fromm landline. Delivery within 28 days."
+    }
+
+    output = predictor.predict(data)
+    assert "score" in output[0]
+
+    # delete predictor
+    predictor.delete_predictor(wait=True)
+
+    # delete endpoint
+    predictor.delete_endpoint()

tests/unit/sagemaker/model/test_model.py

@@ -32,6 +32,8 @@
 from sagemaker.sklearn.model import SKLearnModel
 from sagemaker.tensorflow.model import TensorFlowModel
 from sagemaker.xgboost.model import XGBoostModel
+from sagemaker.enums import EndpointType
+from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
 from sagemaker.workflow.properties import Properties
 from tests.unit import (
     _test_default_bucket_and_prefix_combinations,
@@ -854,6 +856,65 @@ def test_script_mode_model_uses_jumpstart_base_name(repack_model, sagemaker_sess
     ].startswith(JUMPSTART_RESOURCE_BASE_NAME)
 
 
+@patch("sagemaker.utils.repack_model")
+@patch("sagemaker.fw_utils.tar_and_upload_dir")
+def test_all_framework_models_generation_two_endpoint_deploy_path(
+    repack_model, tar_and_uload_dir, sagemaker_session
+):
+    framework_model_classes_to_kwargs = {
+        PyTorchModel: {"framework_version": "1.5.0", "py_version": "py3"},
+        TensorFlowModel: {
+            "framework_version": "2.3",
+        },
+        HuggingFaceModel: {
+            "pytorch_version": "1.7.1",
+            "py_version": "py36",
+            "transformers_version": "4.6.1",
+        },
+        MXNetModel: {"framework_version": "1.7.0", "py_version": "py3"},
+        SKLearnModel: {
+            "framework_version": "0.23-1",
+        },
+        XGBoostModel: {
+            "framework_version": "1.3-1",
+        },
+    }
+
+    sagemaker_session.settings = SessionSettings(include_jumpstart_tags=False)
+
+    source_dir = "s3://blah/blah/blah"
+    for framework_model_class, kwargs in framework_model_classes_to_kwargs.items():
+        framework_model_class(
+            entry_point=ENTRY_POINT_INFERENCE,
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+            model_data=source_dir,
+            **kwargs,
+        ).deploy(
+            instance_type="ml.m2.xlarge",
+            initial_instance_count=INSTANCE_COUNT,
+            endpoint_type=EndpointType.GEN2,
+            resources=ResourceRequirements(
+                requests={
+                    "num_accelerators": 1,
+                    "memory": 8192,
+                    "copies": 1,
+                },
+                limits={},
+            ),
+        )
+
+        # Verified Generation2 endpoint and inference component creation
+        # path
+        sagemaker_session.endpoint_in_service_or_not.assert_called_once()
+        sagemaker_session.create_model.assert_called_once()
+        sagemaker_session.create_inference_component.assert_called_once()
+
+        sagemaker_session.create_inference_component.reset_mock()
+        sagemaker_session.endpoint_in_service_or_not.reset_mock()
+        sagemaker_session.create_model.reset_mock()
+
+
 @patch("sagemaker.utils.repack_model")
 def test_repack_code_location_with_key_prefix(repack_model, sagemaker_session):