infra: use Model class for model deployment unit tests (#1418)

This commit also adds a few more unit tests to round out
coverage for the deploy() method.

Author: laurenyu

tests/unit/sagemaker/model/test_framework_model.py

@@ -24,7 +24,6 @@
 MODEL_DATA = "s3://bucket/model.tar.gz"
 MODEL_IMAGE = "mi"
 ENTRY_POINT = "blah.py"
-INSTANCE_TYPE = "p2.xlarge"
 ROLE = "some-role"
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
@@ -172,145 +171,6 @@ def test_create_no_defaults(sagemaker_session, tmpdir):
     }
 
 
-@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
-@patch("time.strftime", MagicMock(return_value=TIMESTAMP))
-def test_deploy(sagemaker_session, tmpdir):
-    model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir))
-    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1)
-    sagemaker_session.endpoint_from_production_variants.assert_called_with(
-        name=MODEL_NAME,
-        production_variants=[
-            {
-                "InitialVariantWeight": 1,
-                "ModelName": MODEL_NAME,
-                "InstanceType": INSTANCE_TYPE,
-                "InitialInstanceCount": 1,
-                "VariantName": "AllTraffic",
-            }
-        ],
-        tags=None,
-        kms_key=None,
-        wait=True,
-        data_capture_config_dict=None,
-    )
-
-
-@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
-@patch("time.strftime", MagicMock(return_value=TIMESTAMP))
-def test_deploy_endpoint_name(sagemaker_session, tmpdir):
-    model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir))
-    model.deploy(endpoint_name="blah", instance_type=INSTANCE_TYPE, initial_instance_count=55)
-    sagemaker_session.endpoint_from_production_variants.assert_called_with(
-        name="blah",
-        production_variants=[
-            {
-                "InitialVariantWeight": 1,
-                "ModelName": MODEL_NAME,
-                "InstanceType": INSTANCE_TYPE,
-                "InitialInstanceCount": 55,
-                "VariantName": "AllTraffic",
-            }
-        ],
-        tags=None,
-        kms_key=None,
-        wait=True,
-        data_capture_config_dict=None,
-    )
-
-
-@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
-@patch("time.strftime", MagicMock(return_value=TIMESTAMP))
-def test_deploy_tags(sagemaker_session, tmpdir):
-    model = DummyFrameworkModel(sagemaker_session, source_dir=str(tmpdir))
-    tags = [{"ModelName": "TestModel"}]
-    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1, tags=tags)
-    sagemaker_session.endpoint_from_production_variants.assert_called_with(
-        name=MODEL_NAME,
-        production_variants=[
-            {
-                "InitialVariantWeight": 1,
-                "ModelName": MODEL_NAME,
-                "InstanceType": INSTANCE_TYPE,
-                "InitialInstanceCount": 1,
-                "VariantName": "AllTraffic",
-            }
-        ],
-        tags=tags,
-        kms_key=None,
-        wait=True,
-        data_capture_config_dict=None,
-    )
-
-
-@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
-@patch("tarfile.open")
-@patch("time.strftime", return_value=TIMESTAMP)
-def test_deploy_accelerator_type(tfo, time, sagemaker_session):
-    model = DummyFrameworkModel(sagemaker_session)
-    model.deploy(
-        instance_type=INSTANCE_TYPE, initial_instance_count=1, accelerator_type=ACCELERATOR_TYPE
-    )
-    sagemaker_session.endpoint_from_production_variants.assert_called_with(
-        name=MODEL_NAME,
-        production_variants=[
-            {
-                "InitialVariantWeight": 1,
-                "ModelName": MODEL_NAME,
-                "InstanceType": INSTANCE_TYPE,
-                "InitialInstanceCount": 1,
-                "VariantName": "AllTraffic",
-                "AcceleratorType": ACCELERATOR_TYPE,
-            }
-        ],
-        tags=None,
-        kms_key=None,
-        wait=True,
-        data_capture_config_dict=None,
-    )
-
-
-@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
-@patch("tarfile.open")
-@patch("time.strftime", return_value=TIMESTAMP)
-def test_deploy_kms_key(tfo, time, sagemaker_session):
-    key = "some-key-arn"
-    model = DummyFrameworkModel(sagemaker_session)
-    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1, kms_key=key)
-    sagemaker_session.endpoint_from_production_variants.assert_called_with(
-        name=MODEL_NAME,
-        production_variants=[
-            {
-                "InitialVariantWeight": 1,
-                "ModelName": MODEL_NAME,
-                "InstanceType": INSTANCE_TYPE,
-                "InitialInstanceCount": 1,
-                "VariantName": "AllTraffic",
-            }
-        ],
-        tags=None,
-        kms_key=key,
-        wait=True,
-        data_capture_config_dict=None,
-    )
-
-
-@patch("sagemaker.session.Session")
-@patch("sagemaker.local.LocalSession")
-@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
-def test_deploy_creates_correct_session(local_session, session, tmpdir):
-    # We expect a LocalSession when deploying to instance_type = 'local'
-    model = DummyFrameworkModel(sagemaker_session=None, source_dir=str(tmpdir))
-    model.deploy(endpoint_name="blah", instance_type="local", initial_instance_count=1)
-    assert model.sagemaker_session == local_session.return_value
-
-    # We expect a real Session when deploying to instance_type != local/local_gpu
-    model = DummyFrameworkModel(sagemaker_session=None, source_dir=str(tmpdir))
-    model.deploy(
-        endpoint_name="remote_endpoint", instance_type="ml.m4.4xlarge", initial_instance_count=2
-    )
-    assert model.sagemaker_session == session.return_value
-
-
 @patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
 def test_deploy_update_endpoint(sagemaker_session, tmpdir):
     model = DummyFrameworkModel(sagemaker_session, source_dir=tmpdir)