add unit tests

Author: laurenyu

src/sagemaker/model.py

@@ -120,9 +120,9 @@ def _init_sagemaker_session_if_does_not_exist(self, instance_type):
             return
 
         if instance_type in ("local", "local_gpu"):
-            return local.LocalSession()
+            self.sagemaker_session = local.LocalSession()
         else:
-            return session.Session()
+            self.sagemaker_session = session.Session()
 
     def prepare_container_def(
         self, instance_type, accelerator_type=None
@@ -221,7 +221,7 @@ def _compilation_job_config(
             else json.dumps(input_shape),
             "Framework": framework,
         }
-        self._init_sagemaker_session_if_does_not_exist(target_instance_type)
+
         role = self.sagemaker_session.expand_role(role)
         output_model_config = {
             "TargetDevice": target_instance_type,
@@ -336,6 +336,7 @@ def compile(
         framework = framework.upper()
         framework_version = self._get_framework_version() or framework_version
 
+        self._init_sagemaker_session_if_does_not_exist(target_instance_family)
         config = self._compilation_job_config(
             target_instance_family,
             input_shape,
@@ -346,7 +347,6 @@ def compile(
             framework,
             tags,
         )
-        self._init_sagemaker_session_if_does_not_exist(target_instance_family)
         self.sagemaker_session.compile_model(**config)
         job_status = self.sagemaker_session.wait_for_compilation_job(job_name)
         self.model_data = job_status["ModelArtifacts"]["S3ModelArtifacts"]

tests/unit/test_model.py

@@ -435,6 +435,21 @@ def test_model_create_transformer(sagemaker_session):
     sagemaker.model.Model._create_sagemaker_model.assert_called_with(instance_type, tags=tags)
 
 
+@patch("sagemaker.session.Session")
+@patch("sagemaker.local.LocalSession")
+@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
+def test_transformer_creates_correct_session(local_session, session):
+    model = DummyFrameworkModel(sagemaker_session=None)
+    transformer = model.transformer(instance_count=1, instance_type="local")
+    assert model.sagemaker_session == local_session.return_value
+    assert transformer.sagemaker_session == local_session.return_value
+
+    model = DummyFrameworkModel(sagemaker_session=None)
+    transformer = model.transformer(instance_count=1, instance_type="ml.m5.xlarge")
+    assert model.sagemaker_session == session.return_value
+    assert transformer.sagemaker_session == session.return_value
+
+
 def test_model_package_enable_network_isolation_with_no_product_id(sagemaker_session):
     sagemaker_session.sagemaker_client.describe_model_package = Mock(
         return_value=DESCRIBE_MODEL_PACKAGE_RESPONSE
@@ -561,6 +576,21 @@ def test_compile_model_for_cloud(sagemaker_session, tmpdir):
     assert model._is_compiled_model is True
 
 
+@patch("sagemaker.session.Session")
+@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
+def test_compile_creates_session(session):
+    model = DummyFrameworkModel(sagemaker_session=None)
+    model.compile(
+        target_instance_family="ml_c4",
+        input_shape={"data": [1, 3, 1024, 1024]},
+        output_path="s3://output",
+        role="role",
+        framework="tensorflow",
+    )
+
+    assert model.sagemaker_sesion == session.return_value
+
+
 def test_check_neo_region(sagemaker_session, tmpdir):
     sagemaker_session.wait_for_compilation_job = Mock(
         return_value=DESCRIBE_COMPILATION_JOB_RESPONSE