1  import importlib
 2  import importlib.metadata
 3  import re
 4  from typing import Literal
 5  
 6  from packaging.version import InvalidVersion, Version
 7  
 8  from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS, FLAVOR_TO_MODULE_NAME
 9  from mlflow.utils.databricks_utils import is_in_databricks_runtime
10  
11  
12  def _check_version_in_range(ver, min_ver, max_ver):
13      return Version(min_ver) <= Version(ver) <= Version(max_ver)
14  
15  
16  def _check_spark_version_in_range(ver, min_ver, max_ver):
17      """
18      Utility function for allowing late addition release changes to PySpark minor version increments
19      to be accepted, provided that the previous minor version has been previously validated.
20      For example, if version 3.2.1 has been validated as functional with MLflow, an upgrade of
21      PySpark's minor version to 3.2.2 will still provide a valid version check.
22      """
23      parsed_ver = Version(ver)
24      if parsed_ver > Version(min_ver):
25          ver = f"{parsed_ver.major}.{parsed_ver.minor}"
26      return _check_version_in_range(ver, min_ver, max_ver)
27  
28  
29  def _violates_pep_440(ver):
30      try:
31          Version(ver)
32          return False
33      except InvalidVersion:
34          return True
35  
36  
37  def _is_pre_or_dev_release(ver):
38      v = Version(ver)
39      return v.is_devrelease or v.is_prerelease
40  
41  
42  def _strip_dev_version_suffix(version):
43      return re.sub(r"(\.?)dev.*", "", version)
44  
45  
46  def get_min_max_version_and_pip_release(
47      flavor_name: str, category: Literal["autologging", "models"] = "autologging"
48  ):
49      if flavor_name == "pyspark.ml":
50          # pyspark.ml is a special case of spark flavor
51          flavor_name = "spark"
52  
53      min_version = _ML_PACKAGE_VERSIONS[flavor_name][category]["minimum"]
54      max_version = _ML_PACKAGE_VERSIONS[flavor_name][category]["maximum"]
55      pip_release = _ML_PACKAGE_VERSIONS[flavor_name]["package_info"]["pip_release"]
56      return min_version, max_version, pip_release
57  
58  
59  def is_flavor_supported_for_associated_package_versions(flavor_name, check_max_version=True):
60      """
61      Returns:
62          True if the specified flavor is supported for the currently-installed versions of its
63          associated packages.
64      """
65      module_name = FLAVOR_TO_MODULE_NAME[flavor_name]
66  
67      try:
68          actual_version = importlib.import_module(module_name).__version__
69      except AttributeError:
70          try:
71              # NB: Module name is not necessarily the same as the package name. However,
72              # we assume they are the same here for simplicity. If they are not the same,
73              # this will fail and fallback to 'True', which is not a disaster.
74              actual_version = importlib.metadata.version(module_name)
75          except importlib.metadata.PackageNotFoundError:
76              # Some package (e.g. dspy) do not publish version info in a standard format.
77              # For this case, we assume the package version is supported by MLflow.
78              return True
79  
80      # In Databricks, treat 'pyspark 3.x.y.dev0' as 'pyspark 3.x.y'
81      if module_name == "pyspark" and is_in_databricks_runtime():
82          actual_version = _strip_dev_version_suffix(actual_version)
83  
84      if _violates_pep_440(actual_version) or _is_pre_or_dev_release(actual_version):
85          return False
86      min_version, max_version, _ = get_min_max_version_and_pip_release(flavor_name)
87  
88      if module_name == "pyspark" and is_in_databricks_runtime():
89          # MLflow 1.25.0 is known to be compatible with PySpark 3.3.0 on Databricks, despite the
90          # fact that PySpark 3.3.0 was not available in PyPI at the time of the MLflow 1.25.0 release
91          if Version(max_version) < Version("3.3.0"):
92              max_version = "3.3.0"
93          return _check_spark_version_in_range(actual_version, min_version, max_version)
94      else:
95          return (
96              _check_version_in_range(actual_version, min_version, max_version)
97              if check_max_version
98              else Version(min_version) <= Version(actual_version)
99          )