aws / aws/sagemaker-python-sdk

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

Closed
#6,130 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.3k
Forks
1.3k
Avg merge
2d 2h
Merged PRs (30d)
32

Description

**PySDK Version**
- [x] PySDK V2 (2.x)
- [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:

```python
# 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:

```python
# 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:

```python
from sagemaker.jumpstart.model import JumpStartModel

# Any model_id whose specs set inference_vulnerable=true.
model = JumpStartModel(
model_id="",
tolerate_vulnerable_model=True,
tolerate_deprecated_model=True,
)
```

```
sagemaker.jumpstart.exceptions.VulnerableJumpStartModelError: Version '1.0.0' of JumpStart model
'' 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:

```python
>>> 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:

```python
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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.