def test_model_prepare_container_def_no_instance_type_or_image():
model = MXNetModel(MODEL_DATA, role=ROLE, entry_point=SCRIPT_PATH)
with pytest.raises(ValueError) as e:
model.prepare_container_def()
expected_msg = "Must supply either an instance type (for choosing CPU vs GPU) or an image URI."
assert expected_msg in str(e)
def test_model_image_accelerator(sagemaker_session):
model = MXNetModel(MODEL_DATA,
role=ROLE,
entry_point=SCRIPT_PATH,
sagemaker_session=sagemaker_session)
container_def = model.prepare_container_def(
INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE)
assert container_def['Image'] == _get_full_image_uri_with_ei(
defaults.MXNET_VERSION)
def test_model_image_accelerator_mms_version(sagemaker_session):
model = MXNetModel(MODEL_DATA,
role=ROLE,
entry_point=SCRIPT_PATH,
framework_version=MXNetModel._LOWEST_MMS_VERSION,
sagemaker_session=sagemaker_session)
container_def = model.prepare_container_def(
INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE)
assert container_def['Image'] == _get_full_image_uri_with_ei(
MXNetModel._LOWEST_MMS_VERSION, IMAGE_REPO_SERVING_NAME)
def test_model_image_accelerator(
retrieve_image_uri,
repack_model,
tar_and_upload,
sagemaker_session,
mxnet_eia_version,
mxnet_eia_py_version,
):
model = MXNetModel(
MODEL_DATA,
role=ROLE,
entry_point=SCRIPT_PATH,
framework_version=mxnet_eia_version,
py_version=mxnet_eia_py_version,
sagemaker_session=sagemaker_session,
)
container_def = model.prepare_container_def(INSTANCE_TYPE, accelerator_type=ACCELERATOR_TYPE)
assert container_def["Image"] == IMAGE
assert _is_mms_version(mxnet_eia_version) ^ (tar_and_upload.called and not repack_model.called)
