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)