def test_mk_model(self, prj):
model_cfg = cli_utils.mk_model(prj, "model-1", "training-job-1")
assert model_cfg == {
"ModelName": "modelling-project-model-1",
"PrimaryContainer": {
"Image":
"123456789012.dkr.ecr.eu-west-1.amazonaws.com/"
"modelling-project-sagemaker:latest",
"ModelDataUrl":
"s3://prodigyfinance-modelling-project-sagemaker"
"-production/models/modelling-project-training-job-1/"
"output/model.tar.gz",
"Environment": {
"ML2P_MODEL_VERSION":
"modelling-project-model-1",
"ML2P_PROJECT":
"modelling-project",
"ML2P_S3_URL":
("s3://prodigyfinance-modelling-project-sagemaker-production/"
),
},
},
"ExecutionRoleArn":
"arn:aws:iam::111111111111:role/modelling-project",
"Tags": [{
"Key": "ml2p-project",
"Value": "modelling-project"
}],
"EnableNetworkIsolation": False,
}
def test_mk_model_with_vpc_config(self, prj):
prj.deploy["vpc_config"] = {
"security_groups": ["sg-1"],
"subnets": ["net-2"]
}
model_cfg = cli_utils.mk_model(prj, "model-1", "training-job-1")
assert model_cfg["VpcConfig"] == {
"SecurityGroupIds": ["sg-1"],
"Subnets": ["net-2"],
}
def test_mk_model_with_model_type(self, prj):
prj.models["model-type-1"] = "my.pkg.model"
model_cfg = cli_utils.mk_model(prj, "model-1", "training-job-1",
"model-type-1")
assert model_cfg["PrimaryContainer"]["Environment"] == {
"ML2P_MODEL_CLS":
"my.pkg.model",
"ML2P_MODEL_VERSION":
"modelling-project-model-1",
"ML2P_PROJECT":
"modelling-project",
"ML2P_S3_URL":
("s3://prodigyfinance-modelling-project-sagemaker-production/"),
}
def test_mk_model_with_record_invokes(self, prj):
prj.deploy["record_invokes"] = True
model_cfg = cli_utils.mk_model(prj, "model-1", "training-job-1")
assert (model_cfg["PrimaryContainer"]["Environment"]
["ML2P_RECORD_INVOKES"] == "true")
def test_mk_model_with_missing_model_type(self, prj):
with pytest.raises(KeyError) as err:
cli_utils.mk_model(prj, "model-1", "training-job-1",
"model-type-1")
assert str(err.value) == "'model-type-1'"