wehub/mlflow--mlflow
1
0
Fork 0
Code Issues Pull Requests Actions 27 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
mlflow--mlflow/tests/utils/test_requirements_utils.py
T
wehub-resource-sync 4b22cfda96 chore: import upstream snapshot with attribution
2026-07-13 13:22:34 +08:00

771 lines
28 KiB
Python
Raw Permalink Blame History

import importlib
import os
import sys
from importlib.metadata import version
from unittest import mock
import cloudpickle
import importlib_metadata
import pytest
import mlflow
import mlflow.utils.requirements_utils
from mlflow.exceptions import MlflowException
from mlflow.utils.environment import infer_pip_requirements
from mlflow.utils.os import is_windows
from mlflow.utils.requirements_utils import (
_capture_imported_modules,
_check_requirement_satisfied,
_get_installed_version,
_get_pinned_requirement,
_infer_requirements,
_is_comment,
_is_empty,
_is_requirements_file,
_join_continued_lines,
_normalize_package_name,
_parse_requirements,
_prune_packages,
_strip_inline_comment,
_strip_local_version_label,
warn_dependency_requirement_mismatches,
)
from tests.helper_functions import AnyStringWith
def test_is_comment():
assert _is_comment("# comment")
assert _is_comment("#")
assert _is_comment("### comment ###")
assert not _is_comment("comment")
assert not _is_comment("")
def test_is_empty():
assert _is_empty("")
assert not _is_empty(" ")
assert not _is_empty("a")
def test_is_requirements_file():
assert _is_requirements_file("-r req.txt")
assert _is_requirements_file("-r req.txt")
assert _is_requirements_file("--requirement req.txt")
assert _is_requirements_file("--requirement req.txt")
assert not _is_requirements_file("req")
def test_strip_inline_comment():
assert _strip_inline_comment("aaa # comment") == "aaa"
assert _strip_inline_comment("aaa # comment") == "aaa"
assert _strip_inline_comment("aaa # comment") == "aaa"
assert _strip_inline_comment("aaa # com1 # com2") == "aaa"
# Ensure a URI fragment is not stripped
assert (
_strip_inline_comment("git+https://git/repo.git#subdirectory=subdir")
== "git+https://git/repo.git#subdirectory=subdir"
)
def test_join_continued_lines():
assert list(_join_continued_lines(["a"])) == ["a"]
assert list(_join_continued_lines(["a\\", "b"])) == ["ab"]
assert list(_join_continued_lines(["a\\", "b\\", "c"])) == ["abc"]
assert list(_join_continued_lines(["a\\", " b"])) == ["a b"]
assert list(_join_continued_lines(["a\\", " b\\", " c"])) == ["a b c"]
assert list(_join_continued_lines(["a\\", "\\", "b"])) == ["ab"]
assert list(_join_continued_lines(["a\\", "b", "c\\", "d"])) == ["ab", "cd"]
assert list(_join_continued_lines(["a\\", "", "b"])) == ["a", "b"]
assert list(_join_continued_lines(["a\\"])) == ["a"]
assert list(_join_continued_lines(["\\", "a"])) == ["a"]
def test_parse_requirements(tmp_path, monkeypatch):
root_req_src = """
# No version specifier
noverspec
no-ver-spec
# Version specifiers
verspec<1.0
ver-spec == 2.0
# Environment marker
env-marker; python_version < "3.8"
inline-comm # Inline comment
inlinecomm # Inline comment
# Git URIs
git+https://github.com/git/uri
git+https://github.com/sub/dir#subdirectory=subdir
# Requirements files
-r {relative_req}
--requirement {absolute_req}
# Constraints files
-c {relative_con}
--constraint {absolute_con}
# Line continuation
line-cont\
==\
1.0
# Line continuation with spaces
line-cont-space \
== \
1.0
# Line continuation with a blank line
line-cont-blank\
# Line continuation at EOF
line-cont-eof\
""".strip()
monkeypatch.chdir(tmp_path)
root_req = tmp_path.joinpath("requirements.txt")
# Requirements files
rel_req = tmp_path.joinpath("relative_req.txt")
abs_req = tmp_path.joinpath("absolute_req.txt")
# Constraints files
rel_con = tmp_path.joinpath("relative_con.txt")
abs_con = tmp_path.joinpath("absolute_con.txt")
# pip's requirements parser collapses an absolute requirements file path:
# https://github.com/pypa/pip/issues/10121
# As a workaround, use a relative path on Windows.
absolute_req = abs_req.name if is_windows() else str(abs_req)
absolute_con = abs_con.name if is_windows() else str(abs_con)
root_req.write_text(
root_req_src.format(
relative_req=rel_req.name,
absolute_req=absolute_req,
relative_con=rel_con.name,
absolute_con=absolute_con,
)
)
rel_req.write_text("rel-req-xxx\nrel-req-yyy")
abs_req.write_text("abs-req-zzz")
rel_con.write_text("rel-con-xxx\nrel-con-yyy")
abs_con.write_text("abs-con-zzz")
# Uncomment this to get the expected output from pip's internal parser
# from pip._internal.network.session import PipSession
# from pip._internal.req import parse_requirements as pip_parse_requirements
#
# pip_reqs = list(pip_parse_requirements(root_req.name, session=PipSession()))
# print(f"expected_reqs = {[r.requirement for r in pip_reqs if not r.constraint]}")
# print(f"expected_cons = {[r.requirement for r in pip_reqs if r.constraint]}")
expected_reqs = [
"noverspec",
"no-ver-spec",
"verspec<1.0",
"ver-spec == 2.0",
'env-marker; python_version < "3.8"',
"inline-comm",
"inlinecomm",
"git+https://github.com/git/uri",
"git+https://github.com/sub/dir#subdirectory=subdir",
"rel-req-xxx",
"rel-req-yyy",
"abs-req-zzz",
"line-cont==1.0",
"line-cont-space == 1.0",
"line-cont-blank",
"line-cont-eof",
]
expected_cons = [
"rel-con-xxx",
"rel-con-yyy",
"abs-con-zzz",
]
parsed_reqs = list(_parse_requirements(root_req.name, is_constraint=False))
assert [r.req_str for r in parsed_reqs if not r.is_constraint] == expected_reqs
assert [r.req_str for r in parsed_reqs if r.is_constraint] == expected_cons
def test_normalize_package_name():
assert _normalize_package_name("abc") == "abc"
assert _normalize_package_name("ABC") == "abc"
assert _normalize_package_name("a-b-c") == "a-b-c"
assert _normalize_package_name("a.b.c") == "a-b-c"
assert _normalize_package_name("a_b_c") == "a-b-c"
assert _normalize_package_name("a--b--c") == "a-b-c"
assert _normalize_package_name("a-._b-._c") == "a-b-c"
def test_prune_packages():
assert _prune_packages(["mlflow"]) == {"mlflow"}
assert _prune_packages(["mlflow", "scikit-learn"]) == {"mlflow", "scikit-learn"}
def test_capture_imported_modules():
from mlflow.utils._capture_modules import _CaptureImportedModules
with _CaptureImportedModules() as cap:
import math # clint: disable=lazy-import # noqa: F401
__import__("pandas")
importlib.import_module("numpy")
assert "math" in cap.imported_modules
assert "pandas" in cap.imported_modules
assert "numpy" in cap.imported_modules
def test_strip_local_version_label():
assert _strip_local_version_label("1.2.3") == "1.2.3"
assert _strip_local_version_label("1.2.3+ab") == "1.2.3"
assert _strip_local_version_label("1.2.3rc0+ab") == "1.2.3rc0"
assert _strip_local_version_label("1.2.3.dev0+ab") == "1.2.3.dev0"
assert _strip_local_version_label("1.2.3.post0+ab") == "1.2.3.post0"
assert _strip_local_version_label("invalid") == "invalid"
def test_get_installed_version(tmp_path, monkeypatch):
assert _get_installed_version("mlflow") == mlflow.__version__
assert _get_installed_version("numpy") == version("numpy")
assert _get_installed_version("pandas") == version("pandas")
assert _get_installed_version("scikit-learn", module="sklearn") == version("scikit-learn")
not_found_package = tmp_path.joinpath("not_found.py")
not_found_package.write_text("__version__ = '1.2.3'")
monkeypatch.syspath_prepend(str(tmp_path))
with pytest.raises(importlib_metadata.PackageNotFoundError, match=r".+"):
importlib_metadata.version("not_found")
assert _get_installed_version("not_found") == "1.2.3"
def test_package_with_mismatched_pypi_and_import_name():
try:
import dspy # noqa: F401
assert _get_installed_version("dspy") == version("dspy-ai")
except ImportError:
pytest.skip("Skipping test because 'dspy' package is not installed")
def test_get_pinned_requirement(tmp_path, monkeypatch):
assert _get_pinned_requirement("mlflow") == f"mlflow=={mlflow.__version__}"
assert _get_pinned_requirement("mlflow", version="1.2.3") == "mlflow==1.2.3"
not_found_package = tmp_path.joinpath("not_found.py")
not_found_package.write_text("__version__ = '1.2.3'")
monkeypatch.syspath_prepend(str(tmp_path))
with pytest.raises(importlib_metadata.PackageNotFoundError, match=r".+"):
importlib_metadata.version("not_found")
assert _get_pinned_requirement("not_found") == "not_found==1.2.3"
def test_get_pinned_requirement_local_version_label(tmp_path, monkeypatch):
package = tmp_path.joinpath("my_package.py")
lvl = "abc.def.ghi" # Local version label
package.write_text(f"__version__ = '1.2.3+{lvl}'")
monkeypatch.syspath_prepend(str(tmp_path))
with mock.patch("mlflow.utils.requirements_utils._logger.warning") as mock_warning:
req = _get_pinned_requirement("my_package")
mock_warning.assert_called_once()
(first_pos_arg,) = mock_warning.call_args[0]
assert first_pos_arg.startswith(
f"Found my_package version (1.2.3+{lvl}) contains a local version label (+{lvl})."
)
assert req == "my_package==1.2.3"
def test_infer_requirements_excludes_mlflow():
with mock.patch(
"mlflow.utils.requirements_utils._capture_imported_modules",
return_value=["mlflow", "pytest"],
):
mlflow_package = "mlflow-skinny" if "MLFLOW_SKINNY" in os.environ else "mlflow"
assert mlflow_package in importlib_metadata.packages_distributions()["mlflow"]
assert _infer_requirements("path/to/model", "sklearn") == [f"pytest=={pytest.__version__}"]
def test_capture_imported_modules_scopes_databricks_imports(monkeypatch, tmp_path):
from mlflow.utils._capture_modules import _CaptureImportedModules
monkeypatch.chdir(tmp_path)
monkeypatch.syspath_prepend(str(tmp_path))
databricks_dir = os.path.join(tmp_path, "databricks")
os.makedirs(databricks_dir)
for file_name in [
"__init__.py",
"automl.py",
"automl_runtime.py",
"automl_foo.py",
"model_monitoring.py",
"other.py",
]:
with open(os.path.join(databricks_dir, file_name), "w"):
pass
with _CaptureImportedModules() as cap:
# Delete `databricks` from the cache to ensure we load from the dummy module created above.
if "databricks" in sys.modules:
del sys.modules["databricks"]
import databricks
import databricks.automl
import databricks.automl_foo
import databricks.automl_runtime
import databricks.model_monitoring
assert "databricks.automl" in cap.imported_modules
assert "databricks.model_monitoring" in cap.imported_modules
assert "databricks" not in cap.imported_modules
assert "databricks.automl_foo" not in cap.imported_modules
with _CaptureImportedModules() as cap:
import databricks.automl
import databricks.automl_foo
import databricks.automl_runtime
import databricks.model_monitoring
import databricks.other # noqa: F401
assert "databricks.automl" in cap.imported_modules
assert "databricks.model_monitoring" in cap.imported_modules
assert "databricks" in cap.imported_modules
assert "databricks.automl_foo" not in cap.imported_modules
def test_infer_pip_requirements_scopes_databricks_imports():
mlflow.utils.requirements_utils._MODULES_TO_PACKAGES = None
mlflow.utils.requirements_utils._PACKAGES_TO_MODULES = None
with (
mock.patch(
"mlflow.utils.requirements_utils._capture_imported_modules",
return_value=[
"databricks.automl",
"databricks.model_monitoring",
"databricks.automl_runtime",
],
),
mock.patch(
"mlflow.utils.requirements_utils._get_installed_version",
return_value="1.0",
),
mock.patch(
"importlib_metadata.packages_distributions",
return_value={
"databricks": [
"databricks-automl-runtime",
"databricks-model-monitoring",
"koalas",
],
},
),
):
assert _infer_requirements("path/to/model", "sklearn") == [
"databricks-automl-runtime==1.0",
"databricks-model-monitoring==1.0",
]
assert mlflow.utils.requirements_utils._MODULES_TO_PACKAGES["databricks"] == ["koalas"]
def test_capture_imported_modules_include_deps_by_params():
class MyModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
if params is not None:
import pandas as pd
import sklearn # noqa: F401
return pd.DataFrame([params])
return model_input
params = {"a": 1, "b": "string", "c": True}
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
name="test_model",
python_model=MyModel(),
input_example=(["input1"], params),
)
captured_modules = _capture_imported_modules(model_info.model_uri, "pyfunc")
assert "pandas" in captured_modules
assert "sklearn" in captured_modules
@pytest.mark.parametrize(
("module_to_import", "should_capture_extra"),
[
("mlflow.gateway", True),
("mlflow.deployments.server.config", True),
# The `mlflow[gateway]`` extra includes requirements for starting the deployment server,
# but it is not required when the model only uses the deployment client. These test
# cases validate that importing the deployment client alone does not add the extra.
("mlflow.deployments", False),
],
)
def test_capture_imported_modules_includes_gateway_extra(
module_to_import, should_capture_extra, monkeypatch
):
# Disable UV auto-detect to ensure model-based inference is used
monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "false")
class MyModel(mlflow.pyfunc.PythonModel):
def predict(self, context, inputs, params=None):
importlib.import_module(module_to_import)
return inputs
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
name="test_model",
python_model=MyModel(),
input_example=([1, 2, 3]),
)
captured_modules = _capture_imported_modules(model_info.model_uri, "pyfunc")
assert ("mlflow.gateway" in captured_modules) == should_capture_extra
pip_requirements = infer_pip_requirements(model_info.model_uri, "pyfunc")
assert (f"mlflow[gateway]=={mlflow.__version__}" in pip_requirements) == should_capture_extra
def test_gateway_extra_not_captured_when_importing_deployment_client_only(monkeypatch):
# Disable UV auto-detect to ensure model-based inference is used
monkeypatch.setenv("MLFLOW_UV_AUTO_DETECT", "false")
class MyModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
from mlflow.deployments import get_deploy_client # noqa: F401
return model_input
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
name="test_model",
python_model=MyModel(),
input_example=([1, 2, 3]),
)
captured_modules = _capture_imported_modules(model_info.model_uri, "pyfunc")
assert "mlflow.gateway" not in captured_modules
pip_requirements = infer_pip_requirements(model_info.model_uri, "pyfunc")
assert f"mlflow[gateway]=={mlflow.__version__}" not in pip_requirements
def test_warn_dependency_requirement_mismatches():
import sklearn
with mock.patch("mlflow.utils.requirements_utils._logger.warning") as mock_warning:
# Test case: all packages satisfy requirements.
warn_dependency_requirement_mismatches(
model_requirements=[
f"cloudpickle=={cloudpickle.__version__}",
f"scikit-learn=={sklearn.__version__}",
]
)
mock_warning.assert_not_called()
mock_warning.reset_mock()
original_get_installed_version_fn = mlflow.utils.requirements_utils._get_installed_version
def gen_mock_get_installed_version_fn(mock_versions):
def mock_get_installed_version_fn(package, module=None):
if package in mock_versions:
return mock_versions[package]
else:
return original_get_installed_version_fn(package, module)
return mock_get_installed_version_fn
# Test case: multiple mismatched packages
with mock.patch(
"mlflow.utils.requirements_utils._get_installed_version",
gen_mock_get_installed_version_fn({
"scikit-learn": "999.99.11",
"cloudpickle": "999.99.22",
}),
):
warn_dependency_requirement_mismatches(
model_requirements=[
f"cloudpickle=={cloudpickle.__version__}",
f"scikit-learn=={sklearn.__version__}",
]
)
mock_warning.assert_called_once_with(
f"""
Detected one or more mismatches between the model's dependencies and the current Python environment:
- cloudpickle (current: 999.99.22, required: cloudpickle=={cloudpickle.__version__})
- scikit-learn (current: 999.99.11, required: scikit-learn=={sklearn.__version__})
To fix the mismatches, call `mlflow.pyfunc.get_model_dependencies(model_uri)` to fetch the \
model's environment and install dependencies using the resulting environment file.
""".strip()
)
mock_warning.reset_mock()
# Test case: requirement with multiple version specifiers is satisfied
with mock.patch(
"mlflow.utils.requirements_utils._get_installed_version",
gen_mock_get_installed_version_fn({"scikit-learn": "0.8.1"}),
):
warn_dependency_requirement_mismatches(model_requirements=["scikit-learn>=0.8,<=0.9"])
mock_warning.assert_not_called()
mock_warning.reset_mock()
# Test case: requirement with multiple version specifiers is not satisfied
with mock.patch(
"mlflow.utils.requirements_utils._get_installed_version",
gen_mock_get_installed_version_fn({"scikit-learn": "0.7.1"}),
):
warn_dependency_requirement_mismatches(model_requirements=["scikit-learn>=0.8,<=0.9"])
mock_warning.assert_called_once_with(
AnyStringWith(" - scikit-learn (current: 0.7.1, required: scikit-learn>=0.8,<=0.9)")
)
mock_warning.reset_mock()
# Test case: required package is uninstalled.
warn_dependency_requirement_mismatches(model_requirements=["uninstalled-pkg==1.2.3"])
mock_warning.assert_called_once_with(
AnyStringWith(
" - uninstalled-pkg (current: uninstalled, required: uninstalled-pkg==1.2.3)"
)
)
mock_warning.reset_mock()
# Test case: requirement without version specifiers
warn_dependency_requirement_mismatches(model_requirements=["mlflow"])
mock_warning.assert_not_called()
mock_warning.reset_mock()
# Test case: an unexpected error happens while detecting mismatched packages.
with mock.patch(
"mlflow.utils.requirements_utils._check_requirement_satisfied",
side_effect=RuntimeError("check_requirement_satisfied_fn_failed"),
):
warn_dependency_requirement_mismatches(model_requirements=["mlflow"])
mock_warning.assert_called_once_with(
AnyStringWith(
"Encountered an unexpected error "
"(RuntimeError('check_requirement_satisfied_fn_failed')) while "
"detecting model dependency mismatches"
)
)
mock_warning.reset_mock()
# Test case: ignore file path
warn_dependency_requirement_mismatches(model_requirements=["/path/to/my.whl"])
mock_warning.assert_not_called()
def test_check_requirement_satisfied_skips_non_matching_marker():
result = _check_requirement_satisfied("numpy==999.0.0 ; python_full_version < '3.0'")
assert result is None
def test_check_requirement_satisfied_checks_matching_marker():
result = _check_requirement_satisfied("numpy==999.0.0 ; python_full_version >= '3.0'")
assert result is not None
@pytest.mark.parametrize(
"ignore_package_name",
[
"databricks-feature-lookup",
"databricks-agents",
"databricks_agents",
"databricks.agents",
],
)
def test_suppress_warn_dependency_requirement_mismatches_ignore_some_packages(ignore_package_name):
with mock.patch("mlflow.utils.requirements_utils._logger.warning") as mock_warning:
original_get_installed_version_fn = mlflow.utils.requirements_utils._get_installed_version
def gen_mock_get_installed_version_fn(mock_versions):
def mock_get_installed_version_fn(package, module=None):
if package in mock_versions:
return mock_versions[package]
else:
return original_get_installed_version_fn(package, module)
return mock_get_installed_version_fn
# Test case: multiple mismatched packages
with mock.patch(
"mlflow.utils.requirements_utils._get_installed_version",
gen_mock_get_installed_version_fn({
ignore_package_name: "9.99.11",
"cloudpickle": "999.99.22",
}),
):
warn_dependency_requirement_mismatches(
model_requirements=[
f"cloudpickle=={cloudpickle.__version__}",
f"{ignore_package_name}==999.1.1",
]
)
mock_warning.assert_called_once_with(
"""
Detected one or more mismatches between the model's dependencies and the current Python environment:
- cloudpickle (current: 999.99.22, required: cloudpickle=={cloudpickle_version})
To fix the mismatches, call `mlflow.pyfunc.get_model_dependencies(model_uri)` to fetch the \
model's environment and install dependencies using the resulting environment file.
""".strip().format(cloudpickle_version=cloudpickle.__version__)
)
def test_capture_imported_modules_with_exception():
class TestModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
import pandas # noqa: F401
raise Exception("Test exception")
import sklearn # noqa: F401
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
name="model",
python_model=TestModel(),
input_example="test",
)
with mock.patch("mlflow.utils.requirements_utils._logger.warning") as mock_warning:
modules = _capture_imported_modules(model_info.model_uri, mlflow.pyfunc.FLAVOR_NAME)
assert "pandas" in modules
assert (
"Failed to run predict on input_example, dependencies "
"introduced in predict are not captured.\n" in mock_warning.call_args[0][0]
)
assert "sklearn" not in modules
def test_capture_imported_modules_raises_when_env_var_set(monkeypatch):
monkeypatch.setenv("MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS", "True")
class BadModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
raise Exception("Intentional")
with pytest.raises(
MlflowException, match="Encountered an error while capturing imported modules"
):
with mlflow.start_run():
mlflow.pyfunc.log_model(
name="model",
python_model=BadModel(),
input_example="test",
)
def test_capture_imported_modules_correct(monkeypatch):
monkeypatch.setenv("MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS", "true")
class TestModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
import pandas # noqa: F401
import sklearn # noqa: F401
return model_input
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
name="model",
python_model=TestModel(),
input_example="test",
)
modules = _capture_imported_modules(model_info.model_uri, mlflow.pyfunc.FLAVOR_NAME)
assert "pandas" in modules
assert "sklearn" in modules
def test_capture_imported_modules_extra_env_vars(monkeypatch):
monkeypatch.setenv("MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS", "true")
class TestModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
assert os.environ["TEST"] == "test"
return model_input
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
name="model",
python_model=TestModel(),
input_example="test",
pip_requirements=[],
)
_capture_imported_modules(
model_info.model_uri, mlflow.pyfunc.FLAVOR_NAME, extra_env_vars={"TEST": "test"}
)
@pytest.mark.skipif(
importlib.util.find_spec("databricks.agents") is None,
reason="Requires databricks.agents",
)
def test_infer_pip_requirements_on_databricks_agents(tmp_path):
# import here to avoid breaking this test suite on mlflow-skinny
from mlflow.pyfunc import _get_pip_requirements_from_model_path
class TestModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input, params=None):
import databricks.agents # noqa: F401
import pyspark # noqa: F401
return model_input
mlflow.pyfunc.save_model(
tmp_path,
python_model=TestModel(),
input_example="test",
)
requirements = _get_pip_requirements_from_model_path(tmp_path)
packages = [req.split("==")[0] for req in requirements]
assert "databricks-agents" in packages
# databricks-connect should not be pruned even it's a dependency of databricks-agents
assert "databricks-connect" in packages
# pyspark should not exist because it conflicts with databricks-connect
assert "pyspark" not in packages
def test_capture_imported_modules_excludes_pyspark_gateway_env_vars(monkeypatch, tmp_path):
"""
Test that PYSPARK_GATEWAY_PORT and PYSPARK_GATEWAY_SECRET are excluded from the
subprocess environment when capturing imported modules.
These env vars, if inherited by a subprocess, can cause the subprocess to connect
to the parent's py4j gateway. Libraries like databricks-sdk may then corrupt the
parent's gateway state, causing delayed py4j errors like
"Error while obtaining a new communication channel".
"""
monkeypatch.setenv("PYSPARK_GATEWAY_PORT", "12345")
monkeypatch.setenv("PYSPARK_GATEWAY_SECRET", "secret123")
captured_env = {}
def mock_run_command(cmd, timeout_seconds, env):
captured_env.update(env)
raise MlflowException("Mocked - stopping before actual subprocess execution")
with (
mock.patch(
"mlflow.utils.requirements_utils._run_command",
side_effect=mock_run_command,
) as mock_run,
mock.patch(
"mlflow.utils.requirements_utils._download_artifact_from_uri",
return_value=str(tmp_path),
) as mock_download,
):
with pytest.raises(MlflowException, match="Mocked"):
_capture_imported_modules("fake/model/path", "pyfunc")
mock_download.assert_called_once()
mock_run.assert_called_once()
assert "PYSPARK_GATEWAY_PORT" not in captured_env
assert "PYSPARK_GATEWAY_SECRET" not in captured_env
Reference in New Issue View Git Blame Copy Permalink