def test_create_model(sagemaker_session):
source_dir = "s3://mybucket/source"
xgboost_model = XGBoostModel(
model_data=source_dir,
role=ROLE,
sagemaker_session=sagemaker_session,
entry_point=SCRIPT_PATH,
framework_version=XGBOOST_LATEST_VERSION,
)
default_image_uri = _get_full_cpu_image_uri(XGBOOST_LATEST_VERSION)
model_values = xgboost_model.prepare_container_def(CPU)
assert model_values["Image"] == default_image_uri
def test_create_model_with_network_isolation(upload, sagemaker_session,
xgboost_framework_version):
source_dir = "s3://mybucket/source"
repacked_model_data = "s3://mybucket/prefix/model.tar.gz"
xgboost_model = XGBoostModel(
model_data=source_dir,
role=ROLE,
sagemaker_session=sagemaker_session,
entry_point=SCRIPT_PATH,
framework_version=xgboost_framework_version,
enable_network_isolation=True,
)
xgboost_model.uploaded_code = UploadedCode(s3_prefix=repacked_model_data,
script_name="script")
xgboost_model.repacked_model_data = repacked_model_data
model_values = xgboost_model.prepare_container_def(CPU)
assert model_values["Environment"][
"SAGEMAKER_SUBMIT_DIRECTORY"] == "/opt/ml/model/code"
assert model_values["ModelDataUrl"] == repacked_model_data