Remove virtualenv environment if we encounter unexpected error (#8328)

* Remove virtualenv environment if we encounter unexpected error

Signed-off-by: harupy <hkawamura0130@gmail.com>

* Fix message

Signed-off-by: harupy <hkawamura0130@gmail.com>

* test

Signed-off-by: harupy <hkawamura0130@gmail.com>

---------

Signed-off-by: harupy <hkawamura0130@gmail.com>

mlflow/utils/file_utils.py

@@ -11,6 +11,7 @@
 import tempfile
 import stat
 import pathlib
+from contextlib import contextmanager
 
 import urllib.parse
 import urllib.request
@@ -810,3 +811,26 @@ def read_chunk(path: os.PathLike, size: int, start_byte: int = 0) -> bytes:
         if start_byte > 0:
             f.seek(start_byte)
         return f.read(size)
+
+
+@contextmanager
+def remove_on_error(path: os.PathLike, onerror=None):
+    """
+    A context manager that removes a file or directory if an exception is raised during execution.
+
+    :param path: Path to the file or directory.
+    :param onerror: A callback function that will be called with the captured exception before
+                    the file or directory is removed. For example, you can use this callback to
+                    log the exception.
+    """
+    try:
+        yield
+    except Exception as e:
+        if onerror:
+            onerror(e)
+        if os.path.exists(path):
+            if os.path.isfile(path):
+                os.remove(path)
+            elif os.path.isdir(path):
+                shutil.rmtree(path)
+        raise

mlflow/utils/virtualenv.py

@@ -3,13 +3,14 @@
 import shutil
 import uuid
 import re
+import sys
 from pathlib import Path
 from packaging.version import Version
 
 import mlflow
 from mlflow.exceptions import MlflowException
 from mlflow.models.model import Model, MLMODEL_FILE_NAME
-from mlflow.utils.file_utils import TempDir
+from mlflow.utils.file_utils import TempDir, remove_on_error
 from mlflow.utils.process import _exec_cmd, _join_commands, _IS_UNIX
 from mlflow.utils.requirements_utils import _parse_requirements
 from mlflow.utils.environment import (
@@ -236,37 +237,53 @@ def _create_virtualenv(
         _logger.info("Environment %s already exists", env_dir)
         return activate_cmd
 
-    _logger.info("Creating a new environment in %s with %s", env_dir, python_bin_path)
-    _exec_cmd(["virtualenv", "--python", python_bin_path, env_dir], capture_output=capture_output)
-
-    _logger.info("Installing dependencies")
-    for deps in filter(None, [python_env.build_dependencies, python_env.dependencies]):
-        with TempDir() as t:
-            # Create a temporary requirements file in the model directory to resolve the references
-            # in it correctly. To do this, we must first symlink or copy the model directory's
-            # contents to a temporary location for compatibility with deployment tools that store
-            # models in a read-only mount
-            tmp_model_dir = t.path("model")
-            os.makedirs(tmp_model_dir)
-            try:
-                for model_item in os.listdir(local_model_path):
-                    os.symlink(
-                        src=os.path.join(local_model_path, model_item),
-                        dst=os.path.join(tmp_model_dir, model_item),
+    with remove_on_error(
+        env_dir,
+        onerror=lambda e: _logger.warning(
+            "Encountered an unexpected error: %s while creating a virtualenv environment in %s, "
+            "removing the environment directory...",
+            repr(e),
+            env_dir,
+        ),
+    ):
+        _logger.info("Creating a new environment in %s with %s", env_dir, python_bin_path)
+        _exec_cmd(
+            [sys.executable, "-m", "virtualenv", "--python", python_bin_path, env_dir],
+            capture_output=capture_output,
+        )
+
+        _logger.info("Installing dependencies")
+        for deps in filter(None, [python_env.build_dependencies, python_env.dependencies]):
+            with TempDir() as t:
+                # Create a temporary requirements file in the model directory to resolve the
+                # references in it correctly. To do this, we must first symlink or copy the model
+                # directory's contents to a temporary location for compatibility with deployment
+                # tools that store models in a read-only mount
+                tmp_model_dir = t.path("model")
+                os.makedirs(tmp_model_dir)
+                try:
+                    for model_item in os.listdir(local_model_path):
+                        os.symlink(
+                            src=os.path.join(local_model_path, model_item),
+                            dst=os.path.join(tmp_model_dir, model_item),
+                        )
+                except Exception as e:
+                    _logger.warning(
+                        "Failed to symlink model directory during dependency installation"
+                        " Copying instead. Exception: %s",
+                        e,
                     )
-            except Exception as e:
-                _logger.warning(
-                    "Failed to symlink model directory during dependency installation"
-                    " Copying instead. Exception: %s",
-                    e,
-                )
-                shutil.rmtree(tmp_model_dir)
-                _copy_model_to_writeable_destination(local_model_path, tmp_model_dir)
+                    shutil.rmtree(tmp_model_dir)
+                    _copy_model_to_writeable_destination(local_model_path, tmp_model_dir)
 
-            tmp_req_file = f"requirements.{uuid.uuid4().hex}.txt"
-            Path(tmp_model_dir).joinpath(tmp_req_file).write_text("\n".join(deps))
-            cmd = _join_commands(activate_cmd, f"python -m pip install --quiet -r {tmp_req_file}")
-            _exec_cmd(cmd, capture_output=capture_output, cwd=tmp_model_dir, extra_env=extra_env)
+                tmp_req_file = f"requirements.{uuid.uuid4().hex}.txt"
+                Path(tmp_model_dir).joinpath(tmp_req_file).write_text("\n".join(deps))
+                cmd = _join_commands(
+                    activate_cmd, f"python -m pip install --quiet -r {tmp_req_file}"
+                )
+                _exec_cmd(
+                    cmd, capture_output=capture_output, cwd=tmp_model_dir, extra_env=extra_env
+                )
 
     return activate_cmd
 

tests/pyfunc/test_virtualenv.py

@@ -134,6 +134,27 @@ def test_differenet_requirements_create_different_environments(temp_mlflow_env_r
     assert len(list(temp_mlflow_env_root.iterdir())) == 2
 
 
+@use_temp_mlflow_env_root
+def test_environment_directory_is_cleaned_up_when_unexpected_error_occurs(
+    temp_mlflow_env_root, sklearn_model
+):
+    sklearn_req = "scikit-learn==999.999.999"
+    with mlflow.start_run():
+        model_info1 = mlflow.sklearn.log_model(
+            sklearn_model.model,
+            artifact_path="model",
+            pip_requirements=[sklearn_req],
+        )
+
+    try:
+        serve_and_score(model_info1.model_uri, sklearn_model.X_pred)
+    except Exception:
+        pass
+    else:
+        assert False, "Should have raised an exception"
+    assert len(list(temp_mlflow_env_root.iterdir())) == 0
+
+
 @use_temp_mlflow_env_root
 def test_python_env_file_does_not_exist(sklearn_model):
     with mlflow.start_run():