1  import os
  2  from unittest import mock
  3  
  4  import pytest
  5  
  6  import mlflow
  7  from mlflow.exceptions import MlflowException
  8  from mlflow.utils.virtualenv import _create_virtualenv
  9  from mlflow.utils.yaml_utils import read_yaml, write_yaml
 10  
 11  from tests.projects.utils import (
 12      TEST_VIRTUALENV_CONDA_PROJECT_DIR,
 13      TEST_VIRTUALENV_NO_PYTHON_ENV,
 14      TEST_VIRTUALENV_PROJECT_DIR,
 15  )
 16  
 17  spy_on_create_virtualenv = mock.patch(
 18      "mlflow.projects.backend.local._create_virtualenv", wraps=_create_virtualenv
 19  )
 20  
 21  
 22  @pytest.fixture(autouse=True, scope="module")
 23  def use_dev_mlflow_for_projects():
 24      mlflow_root = os.path.dirname(os.path.dirname(mlflow.__file__))
 25  
 26      conda_env = read_yaml(TEST_VIRTUALENV_CONDA_PROJECT_DIR, "conda.yaml")
 27      conda_pip_dependencies = next(
 28          item for item in conda_env["dependencies"] if isinstance(item, dict) and "pip" in item
 29      )["pip"]
 30      if "mlflow" in conda_pip_dependencies:
 31          conda_pip_dependencies.remove("mlflow")
 32          conda_pip_dependencies.append(mlflow_root)
 33      write_yaml(TEST_VIRTUALENV_CONDA_PROJECT_DIR, "conda.yaml", conda_env, overwrite=True)
 34  
 35      for proj_dir in (TEST_VIRTUALENV_PROJECT_DIR, TEST_VIRTUALENV_NO_PYTHON_ENV):
 36          virtualenv_requirements_path = os.path.join(proj_dir, "requirements.txt")
 37          with open(virtualenv_requirements_path) as f:
 38              virtualenv_requirements = f.readlines()
 39  
 40          with open(virtualenv_requirements_path, "w") as f:
 41              for line in virtualenv_requirements:
 42                  if line.rstrip("\n") != "mlflow":
 43                      f.write(line)
 44                  else:
 45                      f.write(mlflow_root)
 46                      f.write("\n")
 47  
 48  
 49  @spy_on_create_virtualenv
 50  def test_virtualenv_project_execution_virtualenv(create_virtualenv_spy):
 51      submitted_run = mlflow.projects.run(
 52          TEST_VIRTUALENV_PROJECT_DIR, entry_point="test", env_manager="virtualenv"
 53      )
 54      submitted_run.wait()
 55      create_virtualenv_spy.assert_called_once()
 56  
 57  
 58  @spy_on_create_virtualenv
 59  def test_virtualenv_project_execution_uv(create_virtualenv_spy):
 60      submitted_run = mlflow.projects.run(
 61          TEST_VIRTUALENV_PROJECT_DIR, entry_point="test", env_manager="uv"
 62      )
 63      submitted_run.wait()
 64      create_virtualenv_spy.assert_called_once()
 65  
 66  
 67  @spy_on_create_virtualenv
 68  def test_virtualenv_project_execution_without_env_manager(create_virtualenv_spy):
 69      # python_env project should be executed using virtualenv without explicitly specifying
 70      # env_manager="virtualenv"
 71      submitted_run = mlflow.projects.run(TEST_VIRTUALENV_PROJECT_DIR, entry_point="test")
 72      submitted_run.wait()
 73      create_virtualenv_spy.assert_called_once()
 74  
 75  
 76  @spy_on_create_virtualenv
 77  def test_virtualenv_project_execution_no_python_env(create_virtualenv_spy):
 78      """
 79      When an MLproject file doesn't contain a `python_env` key but python_env.yaml exists,
 80      virtualenv should be used as an environment manager.
 81      """
 82      submitted_run = mlflow.projects.run(TEST_VIRTUALENV_NO_PYTHON_ENV, entry_point="test")
 83      submitted_run.wait()
 84      create_virtualenv_spy.assert_called_once()
 85  
 86  
 87  @spy_on_create_virtualenv
 88  def test_virtualenv_project_execution_local(create_virtualenv_spy):
 89      submitted_run = mlflow.projects.run(
 90          TEST_VIRTUALENV_PROJECT_DIR, entry_point="main", env_manager="local"
 91      )
 92      submitted_run.wait()
 93      create_virtualenv_spy.assert_not_called()
 94  
 95  
 96  @spy_on_create_virtualenv
 97  def test_virtualenv_conda_project_execution(create_virtualenv_spy):
 98      submitted_run = mlflow.projects.run(
 99          TEST_VIRTUALENV_CONDA_PROJECT_DIR, entry_point="test", env_manager="virtualenv"
100      )
101      submitted_run.wait()
102      create_virtualenv_spy.assert_called_once()
103  
104  
105  def test_virtualenv_project_execution_conda():
106      with pytest.raises(MlflowException, match="python_env project cannot be executed using conda"):
107          mlflow.projects.run(TEST_VIRTUALENV_PROJECT_DIR, env_manager="conda")
108  
109  
110  @spy_on_create_virtualenv
111  def test_virtualenv_project_no_env_file(create_virtualenv_spy, tmp_path):
112      """
113      When neither python_env.yaml nor conda.yaml is present, virtualenv should be used as an
114      environment manager.
115      """
116      ml_project_file = tmp_path.joinpath("MLproject")
117      ml_project_file.write_text(
118          """
119  name: test
120  entry_points:
121    main:
122      command: |
123        python test.py
124  """
125      )
126      tmp_path.joinpath("test.py").write_text(
127          """
128  import os
129  
130  assert "VIRTUAL_ENV" in os.environ
131  """
132      )
133      mlflow.projects.run(str(tmp_path))
134      create_virtualenv_spy.assert_called_once()
135  
136  
137  @spy_on_create_virtualenv
138  def test_virtualenv_project_no_mlmodel_file(create_virtualenv_spy, tmp_path):
139      tmp_path.joinpath("test.py").write_text(
140          """
141  import os
142  
143  assert "VIRTUAL_ENV" in os.environ
144  """
145      )
146      mlflow.projects.run(str(tmp_path), entry_point="test.py")
147      create_virtualenv_spy.assert_called_once()