From 7d98b4ca69055ed74e3480b4659aab9e3fdc2afd Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Fri, 31 Jul 2026 16:48:10 +0000 Subject: [PATCH 1/2] fix: forward tolerance flags from get_jumpstart_configs JumpStartModel.__init__ ends with a call to get_jumpstart_configs, which accepted no tolerate_vulnerable_model or tolerate_deprecated_model argument. It called verify_model_region_and_return_specs without them, so the callee fell back to its False defaults and re-ran the model gate. Constructing JumpStartModel(..., tolerate_vulnerable_model=True) for a flagged model raised VulnerableJumpStartModelError anyway, which made the flag unusable: the object could not be built, so deploy was unreachable. Add both parameters to get_jumpstart_configs, default them to False to keep current behavior for existing callers, and forward them to verify_model_region_and_return_specs. Pass the attributes already set earlier in JumpStartModel.__init__, and do the same in JumpStartEstimator.list_training_configs, which had the same gap. --- X-AI-Prompt: Why do vulnerable models get skipped during deploy instead of tolerated? X-AI-Tool: claude-code --- src/sagemaker/jumpstart/estimator.py | 2 + src/sagemaker/jumpstart/model.py | 2 + src/sagemaker/jumpstart/utils.py | 13 +++ tests/unit/sagemaker/jumpstart/test_utils.py | 104 +++++++++++++++++++ 4 files changed, 121 insertions(+) diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py index b8b75eaa83..1a8f7afef3 100644 --- a/src/sagemaker/jumpstart/estimator.py +++ b/src/sagemaker/jumpstart/estimator.py @@ -1206,6 +1206,8 @@ def list_training_configs(self) -> List[JumpStartMetadataConfig]: region=self.region, scope=JumpStartScriptScope.TRAINING, sagemaker_session=self.sagemaker_session, + tolerate_vulnerable_model=self.tolerate_vulnerable_model, + tolerate_deprecated_model=self.tolerate_deprecated_model, ) return list(configs_dict.values()) diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index 2853df2edc..f594500e7a 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/src/sagemaker/jumpstart/model.py @@ -401,6 +401,8 @@ def _validate_model_id_and_type(): sagemaker_session=self.sagemaker_session, model_type=self.model_type, hub_arn=self.hub_arn, + tolerate_vulnerable_model=self.tolerate_vulnerable_model, + tolerate_deprecated_model=self.tolerate_deprecated_model, ) def log_subscription_warning(self) -> None: diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py index 15f9e9b52e..a9f18d381e 100644 --- a/src/sagemaker/jumpstart/utils.py +++ b/src/sagemaker/jumpstart/utils.py @@ -1146,9 +1146,20 @@ def get_jumpstart_configs( scope: enums.JumpStartScriptScope = enums.JumpStartScriptScope.INFERENCE, model_type: enums.JumpStartModelType = enums.JumpStartModelType.OPEN_WEIGHTS, hub_arn: Optional[str] = None, + tolerate_vulnerable_model: bool = False, + tolerate_deprecated_model: bool = False, ) -> Dict[str, JumpStartMetadataConfig]: """Returns metadata configs for the given model ID and region. + Args: + tolerate_vulnerable_model (bool): True if vulnerable versions of model + specifications should be tolerated (exception not raised). If False, raises an + exception if the script used by this version of the model has dependencies with known + security vulnerabilities. (Default: False). + tolerate_deprecated_model (bool): True if deprecated models should be tolerated + (exception not raised). False if these models should raise an exception. + (Default: False). + Raises: ValueError: If the script scope is not supported by JumpStart. """ @@ -1160,6 +1171,8 @@ def get_jumpstart_configs( scope=scope, model_type=model_type, hub_arn=hub_arn, + tolerate_vulnerable_model=tolerate_vulnerable_model, + tolerate_deprecated_model=tolerate_deprecated_model, ) if scope == enums.JumpStartScriptScope.INFERENCE: diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py index de9be1d51d..398327491a 100644 --- a/tests/unit/sagemaker/jumpstart/test_utils.py +++ b/tests/unit/sagemaker/jumpstart/test_utils.py @@ -1812,6 +1812,110 @@ def test_get_jumpstart_configs_success( "ml.p3.2xlarge", ] + @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs") + def test_get_jumpstart_configs_vulnerable_model_raises_by_default( + self, + patched_get_model_specs, + ): + def make_vulnerable_spec(*largs, **kwargs): + spec = get_base_spec_with_prototype_configs() + spec.inference_vulnerable = True + spec.inference_vulnerabilities = ["CVE-2024-11393"] + return spec + + patched_get_model_specs.side_effect = make_vulnerable_spec + + with pytest.raises(VulnerableJumpStartModelError): + utils.get_jumpstart_configs( + "mock-region", "mock-model", "mock-model-version", config_names=["gpu-inference"] + ) + + @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs") + def test_get_jumpstart_configs_tolerates_vulnerable_model( + self, + patched_get_model_specs, + ): + def make_vulnerable_spec(*largs, **kwargs): + spec = get_base_spec_with_prototype_configs() + spec.inference_vulnerable = True + spec.inference_vulnerabilities = ["CVE-2024-11393"] + return spec + + patched_get_model_specs.side_effect = make_vulnerable_spec + + configs = utils.get_jumpstart_configs( + "mock-region", + "mock-model", + "mock-model-version", + config_names=["gpu-inference"], + tolerate_vulnerable_model=True, + ) + + assert configs.keys() == {"gpu-inference"} + + @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs") + def test_get_jumpstart_configs_deprecated_model_raises_by_default( + self, + patched_get_model_specs, + ): + def make_deprecated_spec(*largs, **kwargs): + spec = get_base_spec_with_prototype_configs() + spec.deprecated = True + return spec + + patched_get_model_specs.side_effect = make_deprecated_spec + + with pytest.raises(DeprecatedJumpStartModelError): + utils.get_jumpstart_configs( + "mock-region", "mock-model", "mock-model-version", config_names=["gpu-inference"] + ) + + @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs") + def test_get_jumpstart_configs_tolerates_deprecated_model( + self, + patched_get_model_specs, + ): + def make_deprecated_spec(*largs, **kwargs): + spec = get_base_spec_with_prototype_configs() + spec.deprecated = True + return spec + + patched_get_model_specs.side_effect = make_deprecated_spec + + configs = utils.get_jumpstart_configs( + "mock-region", + "mock-model", + "mock-model-version", + config_names=["gpu-inference"], + tolerate_deprecated_model=True, + ) + + assert configs.keys() == {"gpu-inference"} + + @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs") + def test_get_jumpstart_configs_tolerates_vulnerable_training_model( + self, + patched_get_model_specs, + ): + def make_vulnerable_spec(*largs, **kwargs): + spec = get_base_spec_with_prototype_configs() + spec.training_vulnerable = True + spec.training_vulnerabilities = ["CVE-2024-11393"] + return spec + + patched_get_model_specs.side_effect = make_vulnerable_spec + + configs = utils.get_jumpstart_configs( + "mock-region", + "mock-model", + "mock-model-version", + config_names=["gpu-training"], + scope=JumpStartScriptScope.TRAINING, + tolerate_vulnerable_model=True, + ) + + assert configs.keys() == {"gpu-training"} + class TestBenchmarkStats: @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs") From fad5a142a1e58e010d7ecbf689350542c306f7f0 Mon Sep 17 00:00:00 2001 From: Evan Kravitz Date: Fri, 31 Jul 2026 18:05:46 +0000 Subject: [PATCH 2/2] test: cover tolerance forwarding in list_training_configs JumpStartEstimator.list_training_configs had no test, so the two forwarded tolerance arguments were uncovered and codecov reported a project coverage drop. --- X-AI-Prompt: Can you cover the estimator path so codecov stops complaining? X-AI-Tool: claude-code --- .../jumpstart/estimator/test_estimator.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py b/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py index 4a64b413f4..622f57d69e 100644 --- a/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py +++ b/tests/unit/sagemaker/jumpstart/estimator/test_estimator.py @@ -2301,6 +2301,37 @@ def test_estimator_set_config_name( mock_estimator_fit.assert_called_once_with(wait=True, job_name="pt-ic-mobilenet-v2-8675309") + @mock.patch("sagemaker.jumpstart.estimator.get_jumpstart_configs") + @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest") + @mock.patch("sagemaker.jumpstart.factory.estimator.Session") + @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs") + @mock.patch("sagemaker.jumpstart.estimator.Estimator.__init__") + @mock.patch("sagemaker.jumpstart.factory.estimator.JUMPSTART_DEFAULT_REGION_NAME", region) + def test_estimator_list_training_configs_forwards_tolerance( + self, + mock_estimator_init: 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_get_jumpstart_configs.return_value = {} + mock_session.return_value = sagemaker_session + + estimator = JumpStartEstimator( + model_id="pytorch-ic-mobilenet-v2", + tolerate_vulnerable_model=True, + tolerate_deprecated_model=True, + ) + + assert estimator.list_training_configs() == [] + assert mock_get_jumpstart_configs.call_args.kwargs["tolerate_vulnerable_model"] is True + assert mock_get_jumpstart_configs.call_args.kwargs["tolerate_deprecated_model"] is True + @mock.patch("sagemaker.utils.sagemaker_timestamp") @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._get_manifest") @mock.patch("sagemaker.jumpstart.factory.estimator.Session")