def get_header_from_base_header(
_obj: JumpStartModelsCache = None,
region: str = None,
model_id: str = None,
semantic_version_str: str = None,
version: str = None,
) -> JumpStartModelHeader:
if version and semantic_version_str:
raise ValueError(
"Cannot specify both `version` and `semantic_version_str` fields.")
if all([
"pytorch" not in model_id,
"tensorflow" not in model_id,
"huggingface" not in model_id,
"mxnet" not in model_id,
"xgboost" not in model_id,
"catboost" not in model_id,
"lightgbm" not in model_id,
"sklearn" not in model_id,
]):
raise KeyError("Bad model ID")
if region is not None and region not in JUMPSTART_REGION_NAME_SET:
raise ValueError(
f"Region name {region} not supported. Please use one of the supported regions in "
f"{JUMPSTART_REGION_NAME_SET}")
spec = copy.deepcopy(BASE_HEADER)
spec["version"] = version or semantic_version_str
spec["model_id"] = model_id
return JumpStartModelHeader(spec)
def test_jumpstart_model_header():
header_dict = {
"model_id":
"tensorflow-ic-imagenet-inception-v3-classification-4",
"version":
"1.0.0",
"min_version":
"2.49.0",
"spec_key":
"community_models_specs/tensorflow-ic-imagenet-inception-v3-classification-4/specs_v1.0.0.json",
}
header1 = JumpStartModelHeader(header_dict)
assert header1.model_id == "tensorflow-ic-imagenet-inception-v3-classification-4"
assert header1.version == "1.0.0"
assert header1.min_version == "2.49.0"
assert (
header1.spec_key ==
"community_models_specs/tensorflow-ic-imagenet-inception-v3-classification-4/specs_v1.0.0.json"
)
assert header1.to_json() == header_dict
header2 = JumpStartModelHeader({
"model_id":
"pytorch-ic-imagenet-inception-v3-classification-4",
"version":
"1.0.0",
"min_version":
"2.49.0",
"spec_key":
"community_models_specs/tensorflow-ic-imagenet-inception-v3-classification-4/specs_v1.0.0.json",
})
assert header1 != header2
header3 = copy.deepcopy(header1)
assert header1 == header3
def get_formatted_manifest(
manifest: List[Dict],
) -> Dict[JumpStartVersionedModelId, JumpStartModelHeader]:
"""Returns formatted manifest dictionary from raw manifest.
Keys are JumpStartVersionedModelId objects, values are
``JumpStartModelHeader`` objects.
"""
manifest_dict = {}
for header in manifest:
header_obj = JumpStartModelHeader(header)
manifest_dict[
JumpStartVersionedModelId(header_obj.model_id, header_obj.version)
] = header_obj
return manifest_dict
def test_get_formatted_manifest():
mock_manifest = [
{
"model_id":
"tensorflow-ic-imagenet-inception-v3-classification-4",
"version":
"1.0.0",
"min_version":
"2.49.0",
"spec_key":
"community_models_specs/tensorflow-ic-imagenet-inception-v3-classification-4/specs_v1.0.0.json",
},
]
assert utils.get_formatted_manifest(mock_manifest) == {
JumpStartVersionedModelId(
"tensorflow-ic-imagenet-inception-v3-classification-4", "1.0.0"):
JumpStartModelHeader(mock_manifest[0])
}
assert utils.get_formatted_manifest([]) == {}
def test_jumpstart_cache_handles_bad_semantic_version_manifest_key_cache():
cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
cache.clear = MagicMock()
cache._model_id_semantic_version_manifest_key_cache = MagicMock()
cache._model_id_semantic_version_manifest_key_cache.get.side_effect = [
JumpStartVersionedModelId(
"tensorflow-ic-imagenet-inception-v3-classification-4", "999.0.0"
),
JumpStartVersionedModelId("tensorflow-ic-imagenet-inception-v3-classification-4", "1.0.0"),
]
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "1.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic-"
"imagenet-inception-v3-classification-4/specs_v1.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4", semantic_version_str="*"
)
cache.clear.assert_called_once()
cache.clear.reset_mock()
cache._model_id_semantic_version_manifest_key_cache.get.side_effect = [
JumpStartVersionedModelId(
"tensorflow-ic-imagenet-inception-v3-classification-4", "999.0.0"
),
JumpStartVersionedModelId(
"tensorflow-ic-imagenet-inception-v3-classification-4", "987.0.0"
),
]
with pytest.raises(KeyError):
cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="*",
)
cache.clear.assert_called_once()
def test_jumpstart_cache_get_header():
cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "2.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic"
"-imagenet-inception-v3-classification-4/specs_v2.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4", semantic_version_str="*"
)
# See if we can make the same query 2 times consecutively
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "2.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic"
"-imagenet-inception-v3-classification-4/specs_v2.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4", semantic_version_str="*"
)
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "2.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic-"
"imagenet-inception-v3-classification-4/specs_v2.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4", semantic_version_str="2.*"
)
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "2.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic-"
"imagenet-inception-v3-classification-4/specs_v2.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="2.0.*",
)
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "2.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic-"
"imagenet-inception-v3-classification-4/specs_v2.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="2.0.0",
)
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "1.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic-"
"imagenet-inception-v3-classification-4/specs_v1.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="1.0.0",
)
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "1.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic-"
"imagenet-inception-v3-classification-4/specs_v1.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4", semantic_version_str="1.*"
)
assert JumpStartModelHeader(
{
"model_id": "tensorflow-ic-imagenet-inception-v3-classification-4",
"version": "1.0.0",
"min_version": "2.49.0",
"spec_key": "community_models_specs/tensorflow-ic-"
"imagenet-inception-v3-classification-4/specs_v1.0.0.json",
}
) == cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="1.0.*",
)
with pytest.raises(KeyError) as e:
cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="3.*",
)
assert (
"Unable to find model manifest for 'tensorflow-ic-imagenet-inception-v3-classification-4' "
"with version '3.*' compatible with your SageMaker version ('2.68.3'). Consider upgrading "
"your SageMaker library to at least version '4.49.0' so you can use version '3.0.0' of "
"'tensorflow-ic-imagenet-inception-v3-classification-4'." in str(e.value)
)
with pytest.raises(KeyError) as e:
cache.get_header(
model_id="pytorch-ic-imagenet-inception-v3-classification-4", semantic_version_str="3.*"
)
assert (
"Unable to find model manifest for 'pytorch-ic-imagenet-inception-v3-classification-4' with "
"version '3.*'. Visit https://sagemaker.readthedocs.io/en/stable/doc_utils/jumpstart.html "
"for updated list of models. Consider using model ID 'pytorch-ic-imagenet-inception-v3-"
"classification-4' with version '2.0.0'."
) in str(e.value)
with pytest.raises(KeyError) as e:
cache.get_header(model_id="pytorch-ic-", semantic_version_str="*")
assert (
"Unable to find model manifest for 'pytorch-ic-' with version '*'. "
"Visit https://sagemaker.readthedocs.io/en/stable/doc_utils/jumpstart.html "
"for updated list of models. "
"Did you mean to use model ID 'pytorch-ic-imagenet-inception-v3-classification-4'?"
) in str(e.value)
with pytest.raises(KeyError) as e:
cache.get_header(model_id="tensorflow-ic-", semantic_version_str="*")
assert (
"Unable to find model manifest for 'tensorflow-ic-' with version '*'. "
"Visit https://sagemaker.readthedocs.io/en/stable/doc_utils/jumpstart.html "
"for updated list of models. "
"Did you mean to use model ID 'tensorflow-ic-imagenet-inception-"
"v3-classification-4'?"
) in str(e.value)
with pytest.raises(KeyError):
cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="BAD",
)
with pytest.raises(KeyError):
cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="2.1.*",
)
with pytest.raises(KeyError):
cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="3.9.*",
)
with pytest.raises(KeyError):
cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4",
semantic_version_str="5.*",
)
with pytest.raises(KeyError):
cache.get_header(
model_id="tensorflow-ic-imagenet-inception-v3-classification-4-bak",
semantic_version_str="*",
)