Skip to content

JumpStartModel ignores tolerate_vulnerable_model and tolerate_deprecated_model: get_jumpstart_configs does not forward the flags #6130

Description

@evakravi

PySDK Version

  • PySDK V2 (2.x)
  • PySDK V3 (3.x)

Describe the bug

get_jumpstart_configs accepts no tolerate_vulnerable_model or tolerate_deprecated_model argument. It calls verify_model_region_and_return_specs without them, so the callee falls back to its False defaults and re-runs the model gate:

# sagemaker/jumpstart/utils.py
def get_jumpstart_configs(
    region: str,
    model_id: str,
    model_version: str,
    config_names: Optional[List[str]] = None,
    sagemaker_session: Optional[Session] = constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
    scope: enums.JumpStartScriptScope = enums.JumpStartScriptScope.INFERENCE,
    model_type: enums.JumpStartModelType = enums.JumpStartModelType.OPEN_WEIGHTS,
    hub_arn: Optional[str] = None,
) -> Dict[str, JumpStartMetadataConfig]:
    model_specs = verify_model_region_and_return_specs(
        region=region,
        model_id=model_id,
        version=model_version,
        sagemaker_session=sagemaker_session,
        scope=scope,
        model_type=model_type,
        hub_arn=hub_arn,
    )   # <-- tolerate_vulnerable_model / tolerate_deprecated_model not forwarded

JumpStartModel.__init__ calls it as its final statement, after super().__init__() has already succeeded:

# sagemaker/jumpstart/model.py
self._metadata_configs = get_jumpstart_configs(
    region=self.region,
    model_id=self.model_id,
    model_version=self.model_version,
    sagemaker_session=self.sagemaker_session,
    model_type=self.model_type,
    hub_arn=self.hub_arn,
)

So JumpStartModel(..., tolerate_vulnerable_model=True) raises VulnerableJumpStartModelError anyway. The constructor honors the flag through every other lookup and then discards it on the last line. The same applies to tolerate_deprecated_model and DeprecatedJumpStartModelError.

The flags are unusable for any flagged model: the object cannot be constructed, so deploy is unreachable. The failure is also non-obvious, because the traceback points at utils.py inside the gate rather than at anything the caller passed.

To reproduce

Any model whose specs carry inference_vulnerable: true reproduces this. tolerate_vulnerable_model=True is ignored:

from sagemaker.jumpstart.model import JumpStartModel

# Any model_id whose specs set inference_vulnerable=true.
model = JumpStartModel(
    model_id="<vulnerable-model-id>",
    tolerate_vulnerable_model=True,
    tolerate_deprecated_model=True,
)
sagemaker.jumpstart.exceptions.VulnerableJumpStartModelError: Version '1.0.0' of JumpStart model
'<vulnerable-model-id>' has at least 1 vulnerable dependency in the inference script.
...
List of vulnerabilities: CVE-2024-11393
  /site-packages/sagemaker/jumpstart/utils.py:726: VulnerableJumpStartModelError

The dropped argument is visible without a flagged model:

>>> import inspect
>>> from sagemaker.jumpstart.utils import get_jumpstart_configs
>>> [p for p in inspect.signature(get_jumpstart_configs).parameters if "tolerate" in p]
[]

Expected behavior

tolerate_vulnerable_model=True and tolerate_deprecated_model=True suppress the gate through the whole constructor, including the metadata config lookup. Construction completes and deploy proceeds.

Empty configs are a reasonable result for a flagged model. _metadata_configs feeds only list_deployment_configs and the benchmark metrics helpers, so no other code path degrades.

Suggested fix: add the two parameters to get_jumpstart_configs, default both to False to preserve current behavior for existing callers, and forward them to verify_model_region_and_return_specs. Then pass self.tolerate_vulnerable_model and self.tolerate_deprecated_model at the JumpStartModel.__init__ call site. Both attributes are already set earlier in that constructor, so no new plumbing is needed.

JumpStartEstimator.list_training_configs has the same gap and would be fixed by the same signature change.

Screenshots or logs

Included above.

System information

  • SageMaker Python SDK version: 2.257.1 and 2.257.5 confirmed; the same code is present on master at sagemaker-core/src/sagemaker/core/jumpstart/utils.py, so 3.x is affected
  • Framework name (eg. PyTorch) or algorithm (eg. KMeans): JumpStart (JumpStartModel, JumpStartEstimator)
  • Framework version: n/a
  • Python version: 3.10
  • CPU or GPU: both
  • Custom Docker image (Y/N): N

Additional context

In 3.x the same get_jumpstart_configs has no tolerance parameters, and ModelBuilder._ensure_metadata_configs in sagemaker-serve/src/sagemaker/serve/model_builder_utils.py calls it without them, so a flagged model raises there too.

This blocked a release pipeline that deploys JumpStart models for validation. Such a pipeline must deploy a model that is flagged, precisely because it needs to test the model, which is what the tolerance flags are for.

The only workaround we found is to replace the lookup and swallow the two gate errors. It is fragile, because it depends on SDK internals and on which modules bound the name at import time:

import sagemaker.jumpstart.estimator
import sagemaker.jumpstart.model
import sagemaker.jumpstart.utils
from sagemaker.jumpstart.exceptions import DeprecatedJumpStartModelError, VulnerableJumpStartModelError

_original = sagemaker.jumpstart.utils.get_jumpstart_configs


def _tolerant_get_jumpstart_configs(*args, **kwargs):
    try:
        return _original(*args, **kwargs)
    except (VulnerableJumpStartModelError, DeprecatedJumpStartModelError):
        return {}


# Each module that imported the name by value needs the replacement.
for module in (sagemaker.jumpstart.utils, sagemaker.jumpstart.model, sagemaker.jumpstart.estimator):
    module.get_jumpstart_configs = _tolerant_get_jumpstart_configs

Worth noting that in the case that prompted this, the vulnerability flag was itself inherited from an inference script bundle the model never runs. That is a separate metadata concern. The SDK bug is that the documented escape hatch does not work.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions