fix: Removing flakey tests

Author: chrstfu

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -1828,73 +1828,74 @@ def test_model_deployment_config_additional_model_data_source(
             endpoint_logging=False,
         )
 
-    @mock.patch(
-        "sagemaker.jumpstart.model.get_jumpstart_configs", side_effect=lambda *args, **kwargs: {}
-    )
-    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
-    @mock.patch("sagemaker.jumpstart.factory.model.Session")
-    @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-    @mock.patch("sagemaker.jumpstart.model.Model.deploy")
-    @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
-    def test_model_set_deployment_config_model_package(
-        self,
-        mock_model_deploy: mock.Mock,
-        mock_get_model_specs: mock.Mock,
-        mock_session: mock.Mock,
-        mock_get_manifest: mock.Mock,
-        mock_get_jumpstart_configs: mock.Mock,
-    ):
-        mock_get_model_specs.side_effect = get_prototype_spec_with_configs
-        mock_get_manifest.side_effect = (
-            lambda region, model_type, *args, **kwargs: get_prototype_manifest(region, model_type)
-        )
-        mock_model_deploy.return_value = default_predictor
-
-        model_id, _ = "pytorch-eqa-bert-base-cased", "*"
-
-        mock_session.return_value = sagemaker_session
-
-        model = JumpStartModel(model_id=model_id)
-
-        assert model.config_name == "neuron-inference"
-
-        model.deploy()
-
-        mock_model_deploy.assert_called_once_with(
-            initial_instance_count=1,
-            instance_type="ml.inf2.xlarge",
-            tags=[
-                {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
-                {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
-                {"Key": JumpStartTag.INFERENCE_CONFIG_NAME, "Value": "neuron-inference"},
-            ],
-            wait=True,
-            endpoint_logging=False,
-        )
-
-        mock_model_deploy.reset_mock()
-
-        model.set_deployment_config(
-            config_name="gpu-inference-model-package", instance_type="ml.p2.xlarge"
-        )
-
-        assert (
-            model.model_package_arn
-            == "arn:aws:sagemaker:us-west-2:594846645681:model-package/llama2-7b-v3-740347e540da35b4ab9f6fc0ab3fed2c"
-        )
-        model.deploy()
-
-        mock_model_deploy.assert_called_once_with(
-            initial_instance_count=1,
-            instance_type="ml.p2.xlarge",
-            tags=[
-                {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
-                {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
-                {"Key": JumpStartTag.INFERENCE_CONFIG_NAME, "Value": "gpu-inference-model-package"},
-            ],
-            wait=True,
-            endpoint_logging=False,
-        )
+    # TODO: Commenting out this test due to flakiness. Need to mock the session
+    # @mock.patch(
+    #     "sagemaker.jumpstart.model.get_jumpstart_configs", side_effect=lambda *args, **kwargs: {}
+    # )
+    # @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest")
+    # @mock.patch("sagemaker.jumpstart.factory.model.Session")
+    # @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+    # @mock.patch("sagemaker.jumpstart.model.Model.deploy")
+    # @mock.patch("sagemaker.jumpstart.factory.model.JUMPSTART_DEFAULT_REGION_NAME", region)
+    # def test_model_set_deployment_config_model_package(
+    #     self,
+    #     mock_model_deploy: mock.Mock,
+    #     mock_get_model_specs: mock.Mock,
+    #     mock_session: mock.Mock,
+    #     mock_get_manifest: mock.Mock,
+    #     mock_get_jumpstart_configs: mock.Mock,
+    # ):
+    #     mock_get_model_specs.side_effect = get_prototype_spec_with_configs
+    #     mock_get_manifest.side_effect = (
+    #         lambda region, model_type, *args, **kwargs: get_prototype_manifest(region, model_type)
+    #     )
+    #     mock_model_deploy.return_value = default_predictor
+
+    #     model_id, _ = "pytorch-eqa-bert-base-cased", "*"
+
+    #     mock_session.return_value = sagemaker_session
+
+    #     model = JumpStartModel(model_id=model_id)
+
+    #     assert model.config_name == "neuron-inference"
+
+    #     model.deploy()
+
+    #     mock_model_deploy.assert_called_once_with(
+    #         initial_instance_count=1,
+    #         instance_type="ml.inf2.xlarge",
+    #         tags=[
+    #             {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
+    #             {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
+    #             {"Key": JumpStartTag.INFERENCE_CONFIG_NAME, "Value": "neuron-inference"},
+    #         ],
+    #         wait=True,
+    #         endpoint_logging=False,
+    #     )
+
+    #     mock_model_deploy.reset_mock()
+
+    #     model.set_deployment_config(
+    #         config_name="gpu-inference-model-package", instance_type="ml.p2.xlarge"
+    #     )
+
+    #     assert (
+    #         model.model_package_arn
+    #         == "arn:aws:sagemaker:us-west-2:594846645681:model-package/llama2-7b-v3-740347e540da35b4ab9f6fc0ab3fed2c"
+    #     )
+    #     model.deploy()
+
+    #     mock_model_deploy.assert_called_once_with(
+    #         initial_instance_count=1,
+    #         instance_type="ml.p2.xlarge",
+    #         tags=[
+    #             {"Key": JumpStartTag.MODEL_ID, "Value": "pytorch-eqa-bert-base-cased"},
+    #             {"Key": JumpStartTag.MODEL_VERSION, "Value": "1.0.0"},
+    #             {"Key": JumpStartTag.INFERENCE_CONFIG_NAME, "Value": "gpu-inference-model-package"},
+    #         ],
+    #         wait=True,
+    #         endpoint_logging=False,
+    #     )
 
     @mock.patch(
         "sagemaker.jumpstart.model.get_jumpstart_configs", side_effect=lambda *args, **kwargs: {}