/ mlflow / mismatch.py
mismatch.py
 1  from __future__ import annotations
 2  
 3  import importlib.metadata
 4  import warnings
 5  
 6  
 7  def _get_version(package_name: str) -> str | None:
 8      try:
 9          return importlib.metadata.version(package_name)
10      except importlib.metadata.PackageNotFoundError:
11          return None
12  
13  
14  def _check_version_mismatch() -> None:
15      """
16      Warns if both mlflow and child packages are installed but their versions are different.
17  
18      Reference: https://github.com/pypa/pip/issues/4625
19      """
20      mlflow_ver = _get_version("mlflow")
21      # Skip if mlflow is installed from source.
22      if mlflow_ver is None or "dev" in mlflow_ver:
23          return
24  
25      child_packages = ["mlflow-skinny", "mlflow-tracing"]
26      child_versions = [(p, _get_version(p)) for p in child_packages]
27  
28      mismatched = [
29          (p, v) for p, v in child_versions if v is not None and "dev" not in v and v != mlflow_ver
30      ]
31  
32      if mismatched:
33          mismatched_str = ", ".join(f"{name} ({ver})" for name, ver in mismatched)
34          warnings.warn(
35              (
36                  f"Versions of mlflow ({mlflow_ver}) and child packages {mismatched_str} "
37                  "are different. This may lead to unexpected behavior. "
38                  "Please install the same version of all MLflow packages."
39              ),
40              stacklevel=2,
41              category=UserWarning,
42          )