def test_git_support_succeed_model_class(tar_and_upload_dir, git_clone_repo,
sagemaker_session):
git_clone_repo.side_effect = lambda gitconfig, entrypoint, sourcedir, dependency: {
"entry_point": "entry_point",
"source_dir": "/tmp/repo_dir/source_dir",
"dependencies": ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"],
}
entry_point = "entry_point"
source_dir = "source_dir"
dependencies = ["foo", "bar"]
git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
model = Model(
sagemaker_session=sagemaker_session,
entry_point=entry_point,
source_dir=source_dir,
dependencies=dependencies,
git_config=git_config,
image_uri=IMAGE_URI,
)
model.prepare_container_def(instance_type=INSTANCE_TYPE)
git_clone_repo.assert_called_with(git_config, entry_point, source_dir,
dependencies)
assert model.entry_point == "entry_point"
assert model.source_dir == "/tmp/repo_dir/source_dir"
assert model.dependencies == ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"]
def test_prepare_container_def_with_model_data_and_env():
env = {"FOO": "BAR"}
model = Model(MODEL_IMAGE, MODEL_DATA, env=env)
expected = {"Image": MODEL_IMAGE, "Environment": env, "ModelDataUrl": MODEL_DATA}
container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
assert expected == container_def
container_def = model.prepare_container_def()
assert expected == container_def
def test_prepare_container_def_with_image_config():
image_config = {"RepositoryAccessMode": "Vpc"}
model = Model(MODEL_IMAGE, image_config=image_config)
expected = {
"Image": MODEL_IMAGE,
"ImageConfig": {"RepositoryAccessMode": "Vpc"},
"Environment": {},
}
container_def = model.prepare_container_def()
assert expected == container_def
def test_prepare_container_def_with_model_src_s3_returns_correct_url(
sagemaker_session):
model = Model(
entry_point=ENTRY_POINT_INFERENCE,
role=ROLE,
sagemaker_session=sagemaker_session,
source_dir=SCRIPT_URI,
image_uri=MODEL_IMAGE,
model_data=Properties("Steps.MyStep"),
)
container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
assert container_def["Environment"][
"SAGEMAKER_SUBMIT_DIRECTORY"] == SCRIPT_URI
def test_prepare_container_def_with_model_data():
model = Model(MODEL_IMAGE)
container_def = model.prepare_container_def(INSTANCE_TYPE, "ml.eia.medium")
expected = {"Image": MODEL_IMAGE, "Environment": {}}
assert expected == container_def
