Files
2026-07-13 13:22:34 +08:00

128 lines
4.3 KiB
Python

import os
import textwrap
import pytest
from mlflow.exceptions import ExecutionException
from mlflow.projects import _project_spec
from tests.projects.utils import load_project
def test_project_get_entry_point():
project = load_project()
entry_point = project.get_entry_point("greeter")
assert entry_point.name == "greeter"
assert entry_point.command == "python greeter.py {greeting} {name}"
# Validate parameters
assert set(entry_point.parameters.keys()) == {"name", "greeting"}
name_param = entry_point.parameters["name"]
assert name_param.type == "string"
assert name_param.default is None
greeting_param = entry_point.parameters["greeting"]
assert greeting_param.type == "string"
assert greeting_param.default == "hi"
def test_project_get_unspecified_entry_point():
project = load_project()
entry_point = project.get_entry_point("my_script.py")
assert entry_point.name == "my_script.py"
assert entry_point.command == "python my_script.py"
assert entry_point.parameters == {}
entry_point = project.get_entry_point("my_script.sh")
assert entry_point.name == "my_script.sh"
assert entry_point.command == "{} my_script.sh".format(os.environ.get("SHELL", "bash"))
assert entry_point.parameters == {}
with pytest.raises(ExecutionException, match="Could not find my_program.scala"):
project.get_entry_point("my_program.scala")
@pytest.mark.parametrize(
(
# Contents of MLproject file. If None, no MLproject file will be written.
"mlproject",
# Path to conda environment file. If None, no conda environment file will be written.
"conda_env_path",
# Contents of conda environment file (written if conda_env_path is not None).
"conda_env_contents",
# Path to MLproject file. If None, the MLproject file will be written to "MLproject".
"mlproject_path",
),
[
(None, None, "", None),
("key: value", "conda.yaml", "hi", "MLproject"),
("conda_env: some-env.yaml", "some-env.yaml", "hi", "mlproject"),
],
)
def test_load_project(tmp_path, mlproject, conda_env_path, conda_env_contents, mlproject_path):
"""
Test that we can load a project with various combinations of an MLproject / conda.yaml file
"""
if mlproject:
tmp_path.joinpath(mlproject_path).write_text(mlproject)
if conda_env_path:
tmp_path.joinpath(conda_env_path).write_text(conda_env_contents)
project = _project_spec.load_project(str(tmp_path))
assert project._entry_points == {}
expected_env_path = str(tmp_path.joinpath(conda_env_path)) if conda_env_path else None
assert project.env_config_path == expected_env_path
if conda_env_path:
with open(project.env_config_path) as f:
assert f.read() == conda_env_contents
def test_load_docker_project(tmp_path):
tmp_path.joinpath("MLproject").write_text(
textwrap.dedent(
"""
docker_env:
image: some-image
"""
)
)
project = _project_spec.load_project(str(tmp_path))
assert project._entry_points == {}
assert project.env_config_path is None
assert project.docker_env.get("image") == "some-image"
def test_load_virtualenv_project(tmp_path):
tmp_path.joinpath("MLproject").write_text("python_env: python_env.yaml")
python_env = tmp_path.joinpath("python_env.yaml")
python_env.write_text("python: 3.8.15")
project = _project_spec.load_project(tmp_path)
assert project._entry_points == {}
assert python_env.samefile(project.env_config_path)
@pytest.mark.parametrize(
("invalid_project_contents", "expected_error_msg"),
[
(
textwrap.dedent(
"""
docker_env:
image: some-image
conda_env: some-file.yaml
"""
),
"cannot contain multiple environment fields",
),
(
textwrap.dedent(
"""
docker_env:
not-image-attribute: blah
"""
),
"no image attribute found",
),
],
)
def test_load_invalid_project(tmp_path, invalid_project_contents, expected_error_msg):
tmp_path.joinpath("MLproject").write_text(invalid_project_contents)
with pytest.raises(ExecutionException, match=expected_error_msg) as e:
_project_spec.load_project(str(tmp_path))
assert expected_error_msg in str(e.value)