wehub-resource-sync
148 lines
4.7 KiB
Python
148 lines
4.7 KiB
Python
import os
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.utils.virtualenv import _create_virtualenv
|
|
from mlflow.utils.yaml_utils import read_yaml, write_yaml
|
|
|
|
from tests.projects.utils import (
|
|
TEST_VIRTUALENV_CONDA_PROJECT_DIR,
|
|
TEST_VIRTUALENV_NO_PYTHON_ENV,
|
|
TEST_VIRTUALENV_PROJECT_DIR,
|
|
)
|
|
|
|
spy_on_create_virtualenv = mock.patch(
|
|
"mlflow.projects.backend.local._create_virtualenv", wraps=_create_virtualenv
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="module")
|
|
def use_dev_mlflow_for_projects():
|
|
mlflow_root = os.path.dirname(os.path.dirname(mlflow.__file__))
|
|
|
|
conda_env = read_yaml(TEST_VIRTUALENV_CONDA_PROJECT_DIR, "conda.yaml")
|
|
conda_pip_dependencies = next(
|
|
item for item in conda_env["dependencies"] if isinstance(item, dict) and "pip" in item
|
|
)["pip"]
|
|
if "mlflow" in conda_pip_dependencies:
|
|
conda_pip_dependencies.remove("mlflow")
|
|
conda_pip_dependencies.append(mlflow_root)
|
|
write_yaml(TEST_VIRTUALENV_CONDA_PROJECT_DIR, "conda.yaml", conda_env, overwrite=True)
|
|
|
|
for proj_dir in (TEST_VIRTUALENV_PROJECT_DIR, TEST_VIRTUALENV_NO_PYTHON_ENV):
|
|
virtualenv_requirements_path = os.path.join(proj_dir, "requirements.txt")
|
|
with open(virtualenv_requirements_path) as f:
|
|
virtualenv_requirements = f.readlines()
|
|
|
|
with open(virtualenv_requirements_path, "w") as f:
|
|
for line in virtualenv_requirements:
|
|
if line.rstrip("\n") != "mlflow":
|
|
f.write(line)
|
|
else:
|
|
f.write(mlflow_root)
|
|
f.write("\n")
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_project_execution_virtualenv(create_virtualenv_spy):
|
|
submitted_run = mlflow.projects.run(
|
|
TEST_VIRTUALENV_PROJECT_DIR, entry_point="test", env_manager="virtualenv"
|
|
)
|
|
submitted_run.wait()
|
|
create_virtualenv_spy.assert_called_once()
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_project_execution_uv(create_virtualenv_spy):
|
|
submitted_run = mlflow.projects.run(
|
|
TEST_VIRTUALENV_PROJECT_DIR, entry_point="test", env_manager="uv"
|
|
)
|
|
submitted_run.wait()
|
|
create_virtualenv_spy.assert_called_once()
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_project_execution_without_env_manager(create_virtualenv_spy):
|
|
# python_env project should be executed using virtualenv without explicitly specifying
|
|
# env_manager="virtualenv"
|
|
submitted_run = mlflow.projects.run(TEST_VIRTUALENV_PROJECT_DIR, entry_point="test")
|
|
submitted_run.wait()
|
|
create_virtualenv_spy.assert_called_once()
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_project_execution_no_python_env(create_virtualenv_spy):
|
|
"""
|
|
When an MLproject file doesn't contain a `python_env` key but python_env.yaml exists,
|
|
virtualenv should be used as an environment manager.
|
|
"""
|
|
submitted_run = mlflow.projects.run(TEST_VIRTUALENV_NO_PYTHON_ENV, entry_point="test")
|
|
submitted_run.wait()
|
|
create_virtualenv_spy.assert_called_once()
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_project_execution_local(create_virtualenv_spy):
|
|
submitted_run = mlflow.projects.run(
|
|
TEST_VIRTUALENV_PROJECT_DIR, entry_point="main", env_manager="local"
|
|
)
|
|
submitted_run.wait()
|
|
create_virtualenv_spy.assert_not_called()
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_conda_project_execution(create_virtualenv_spy):
|
|
submitted_run = mlflow.projects.run(
|
|
TEST_VIRTUALENV_CONDA_PROJECT_DIR, entry_point="test", env_manager="virtualenv"
|
|
)
|
|
submitted_run.wait()
|
|
create_virtualenv_spy.assert_called_once()
|
|
|
|
|
|
def test_virtualenv_project_execution_conda():
|
|
with pytest.raises(MlflowException, match="python_env project cannot be executed using conda"):
|
|
mlflow.projects.run(TEST_VIRTUALENV_PROJECT_DIR, env_manager="conda")
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_project_no_env_file(create_virtualenv_spy, tmp_path):
|
|
"""
|
|
When neither python_env.yaml nor conda.yaml is present, virtualenv should be used as an
|
|
environment manager.
|
|
"""
|
|
ml_project_file = tmp_path.joinpath("MLproject")
|
|
ml_project_file.write_text(
|
|
"""
|
|
name: test
|
|
entry_points:
|
|
main:
|
|
command: |
|
|
python test.py
|
|
"""
|
|
)
|
|
tmp_path.joinpath("test.py").write_text(
|
|
"""
|
|
import os
|
|
|
|
assert "VIRTUAL_ENV" in os.environ
|
|
"""
|
|
)
|
|
mlflow.projects.run(str(tmp_path))
|
|
create_virtualenv_spy.assert_called_once()
|
|
|
|
|
|
@spy_on_create_virtualenv
|
|
def test_virtualenv_project_no_mlmodel_file(create_virtualenv_spy, tmp_path):
|
|
tmp_path.joinpath("test.py").write_text(
|
|
"""
|
|
import os
|
|
|
|
assert "VIRTUAL_ENV" in os.environ
|
|
"""
|
|
)
|
|
mlflow.projects.run(str(tmp_path), entry_point="test.py")
|
|
create_virtualenv_spy.assert_called_once()
|