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