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