Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_path_params():
with TempDir() as tmp:
dest_path = tmp.path()
data_file = "s3://path.test/resources/data_file.csv"
defaults = {
"constants": {"type": "uri", "default": "s3://path.test/b1"},
"data": {"type": "path", "default": data_file}
}
entry_point = EntryPoint("entry_point_name", defaults, "command_name script.py")
with mock.patch("mlflow.data.download_uri") as download_uri_mock:
final_1, extra_1 = entry_point.compute_parameters({}, None)
assert (final_1 == {"constants": "s3://path.test/b1", "data": data_file})
assert (extra_1 == {})
assert download_uri_mock.call_count == 0
with mock.patch("mlflow.data.download_uri") as download_uri_mock:
user_2 = {"alpha": 0.001, "constants": "s3://path.test/b_two"}
def test_model_log(xgb_model, model_path):
old_uri = mlflow.get_tracking_uri()
model = xgb_model.model
with TempDir(chdr=True, remove_on_exit=True) as tmp:
for should_start_run in [False, True]:
try:
mlflow.set_tracking_uri("test")
if should_start_run:
mlflow.start_run()
artifact_path = "model"
conda_env = os.path.join(tmp.path(), "conda_env.yaml")
_mlflow_conda_env(conda_env, additional_pip_deps=["xgboost"])
mlflow.xgboost.log_model(
xgb_model=model,
artifact_path=artifact_path,
conda_env=conda_env)
model_uri = "runs:/{run_id}/{artifact_path}".format(
run_id=mlflow.active_run().info.run_id,
def test_build_image_with_relative_model_path_calls_expected_azure_routines(
sklearn_model):
with TempDir(chdr=True):
model_path = "model"
mlflow.sklearn.save_model(sk_model=sklearn_model, path=model_path)
with AzureMLMocks() as aml_mocks:
workspace = get_azure_workspace()
mlflow.azureml.build_image(model_uri=model_path, workspace=workspace)
assert aml_mocks["register_model"].call_count == 1
assert aml_mocks["create_image"].call_count == 1
def test_spark_module_model_save_with_relative_path_and_valid_sample_input_produces_mleap_flavor(
spark_model_iris):
with TempDir(chdr=True) as tmp:
model_path = os.path.basename(tmp.path("model"))
mlflow_model = Model()
sparkm.save_model(spark_model=spark_model_iris.model,
path=model_path,
sample_input=spark_model_iris.spark_df,
mlflow_model=mlflow_model)
assert mleap.FLAVOR_NAME in mlflow_model.flavors
config_path = os.path.join(model_path, "MLmodel")
assert os.path.exists(config_path)
config = Model.load(config_path)
assert mleap.FLAVOR_NAME in config.flavors
def test_cli_build_image_with_relative_model_path_calls_expected_azure_routines(sklearn_model):
with TempDir(chdr=True):
model_path = "model"
mlflow.sklearn.save_model(sk_model=sklearn_model, path=model_path)
with AzureMLMocks() as aml_mocks:
result = CliRunner(env={"LC_ALL": "en_US.UTF-8", "LANG": "en_US.UTF-8"}).invoke(
mlflow.azureml.cli.commands,
[
'build-image',
'-m', model_path,
'-w', 'test_workspace',
'-i', 'image_name',
'-n', 'model_name',
])
assert result.exit_code == 0
assert aml_mocks["register_model"].call_count == 1
def test_gbt():
old_uri = tracking.get_tracking_uri()
with TempDir(chdr=False, remove_on_exit=True) as tmp:
try:
diamonds = tmp.path("diamonds")
artifacts = tmp.path("artifacts")
os.mkdir(diamonds)
os.mkdir(artifacts)
tracking.set_tracking_uri(artifacts)
mlflow.set_experiment("test-experiment")
# Download the diamonds dataset via mlflow run
run(".", entry_point="main", version=None,
parameters={"dest-dir": diamonds},
mode="local", cluster_spec=None, git_username=None, git_password=None,
use_conda=True, storage_dir=None)
# Run the main gbt app via mlflow
submitted_run = run(
"apps/gbt-regression", entry_point="main", version=None,
def test_download_artifacts(local_artifact_repo, dst_path):
artifact_rel_path = "test.txt"
artifact_text = "hello world!"
empty_dir_path = "empty_dir"
with TempDir(chdr=True) as local_dir:
if dst_path:
os.mkdir(dst_path)
artifact_src_path = local_dir.path(artifact_rel_path)
os.mkdir(local_dir.path(empty_dir_path))
with open(artifact_src_path, "w") as f:
f.write(artifact_text)
local_artifact_repo.log_artifacts(local_dir.path())
result = local_artifact_repo.download_artifacts(artifact_path=artifact_rel_path,
dst_path=dst_path)
assert open(result).read() == artifact_text
result = local_artifact_repo.download_artifacts(artifact_path="", dst_path=dst_path)
empty_dir_dst_path = os.path.join(result, empty_dir_path)
assert os.path.isdir(empty_dir_dst_path)
assert len(os.listdir(empty_dir_dst_path)) == 0
def test_build_image_includes_mlflow_home_as_file_dependency_if_specified(
sklearn_model, model_path):
def mock_create_dockerfile(output_path, *args, **kwargs):
# pylint: disable=unused-argument
with open(output_path, "w") as f:
f.write("Dockerfile contents")
mlflow.sklearn.save_model(sk_model=sklearn_model, path=model_path)
with AzureMLMocks() as aml_mocks, TempDir() as tmp,\
mock.patch("mlflow.azureml._create_dockerfile") as create_dockerfile_mock:
create_dockerfile_mock.side_effect = mock_create_dockerfile
# Write a mock `setup.py` file to the mlflow home path so that it will be recognized
# as a viable MLflow source directory during the image build process
mlflow_home = tmp.path()
with open(os.path.join(mlflow_home, "setup.py"), "w") as f:
f.write("setup instructions")
workspace = get_azure_workspace()
mlflow.azureml.build_image(
model_uri=model_path, workspace=workspace, mlflow_home=mlflow_home)
assert len(create_dockerfile_mock.call_args_list) == 1
_, create_dockerfile_kwargs = create_dockerfile_mock.call_args_list[0]
# The path to MLflow that is referenced by the Docker container may differ from the
def test_entry_point_compute_command():
"""
Tests that EntryPoint correctly computes the command to execute in order to run the entry point.
"""
project = load_project()
entry_point = project.get_entry_point("greeter")
with TempDir() as tmp:
storage_dir = tmp.path()
command = entry_point.compute_command({"name": "friend", "excitement": 10}, storage_dir)
assert command == "python greeter.py hi friend --excitement 10"
with pytest.raises(ExecutionException):
entry_point.compute_command({}, storage_dir)
# Test shell escaping
name_value = "friend; echo 'hi'"
command = entry_point.compute_command({"name": name_value}, storage_dir)
assert command == "python greeter.py %s %s" % (shlex_quote("hi"), shlex_quote(name_value))
def _get_flavor_backend(model_uri, **kwargs):
with TempDir() as tmp:
if ModelsArtifactRepository.is_models_uri(model_uri):
underlying_model_uri = ModelsArtifactRepository.get_underlying_uri(model_uri)
else:
underlying_model_uri = model_uri
local_path = _download_artifact_from_uri(posixpath.join(underlying_model_uri, "MLmodel"),
output_path=tmp.path())
model = Model.load(local_path)
flavor_name, flavor_backend = get_flavor_backend(model, **kwargs)
if flavor_backend is None:
raise Exception("No suitable flavor backend was found for the model.")
_logger.info("Selected backend for flavor '%s'", flavor_name)
return flavor_backend