fix: tests

Author: evakravi

src/sagemaker/jumpstart/estimator.py

@@ -608,6 +608,7 @@ def _validate_model_id_and_get_type_hook():
         self.role = estimator_init_kwargs.role
         self.sagemaker_session = estimator_init_kwargs.sagemaker_session
         self._enable_network_isolation = estimator_init_kwargs.enable_network_isolation
+        super(JumpStartEstimator, self).__init__(**estimator_init_kwargs.to_kwargs_dict())
         self._internal_config = JumpStartModelInternalConfig(
             specs=verify_model_region_and_return_specs(
                 region=self.region,
@@ -617,9 +618,10 @@ def _validate_model_id_and_get_type_hook():
                 model_type=self.model_type,
                 scope=JumpStartScriptScope.TRAINING,
                 sagemaker_session=self.sagemaker_session,
+                tolerate_vulnerable_model=self.tolerate_vulnerable_model,
+                tolerate_deprecated_model=self.tolerate_deprecated_model,
             )
         )
-        super(JumpStartEstimator, self).__init__(**estimator_init_kwargs.to_kwargs_dict())
 
     def fit(
         self,

src/sagemaker/jumpstart/model.py

@@ -361,6 +361,9 @@ def _validate_model_id_and_type():
             self.log_subscription_warning()
 
         model_init_kwargs_dict = model_init_kwargs.to_kwargs_dict()
+        super(JumpStartModel, self).__init__(**model_init_kwargs_dict)
+
+        self.model_package_arn = model_init_kwargs.model_package_arn
         self._internal_config = JumpStartModelInternalConfig(
             specs=verify_model_region_and_return_specs(
                 region=self.region,
@@ -370,11 +373,10 @@ def _validate_model_id_and_type():
                 model_type=self.model_type,
                 scope=JumpStartScriptScope.INFERENCE,
                 sagemaker_session=self.sagemaker_session,
+                tolerate_vulnerable_model=self.tolerate_vulnerable_model,
+                tolerate_deprecated_model=self.tolerate_deprecated_model,
             )
         )
-        super(JumpStartModel, self).__init__(**model_init_kwargs_dict)
-
-        self.model_package_arn = model_init_kwargs.model_package_arn
 
     def log_subscription_warning(self) -> None:
         """Log message prompting the customer to subscribe to the proprietary model."""

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -1171,8 +1171,10 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
     @mock.patch("sagemaker.jumpstart.estimator.get_init_kwargs")
     @mock.patch("sagemaker.jumpstart.estimator.Estimator.__init__")
     @mock.patch("sagemaker.jumpstart.estimator.validate_model_id_and_get_type")
+    @mock.patch("sagemaker.jumpstart.estimator.verify_model_region_and_return_specs")
     def test_validate_model_id_and_get_type(
         self,
+        verify_model_region_and_return_specs: mock.Mock,
         mock_validate_model_id_and_get_type: mock.Mock,
         mock_init: mock.Mock,
         mock_get_init_kwargs: mock.Mock,

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

@@ -758,8 +758,10 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
     @mock.patch("sagemaker.jumpstart.model.get_init_kwargs")
     @mock.patch("sagemaker.jumpstart.model.Model.__init__")
     @mock.patch("sagemaker.jumpstart.model.validate_model_id_and_get_type")
+    @mock.patch("sagemaker.jumpstart.model.verify_model_region_and_return_specs")
     def test_validate_model_id_and_get_type(
         self,
+        verify_model_region_and_return_specs: mock.Mock,
         mock_validate_model_id_and_get_type: mock.Mock,
         mock_init: mock.Mock,
         mock_get_init_kwargs: mock.Mock,