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'"