import argparse import logging import re from pathlib import Path from packaging.version import Version _logger = logging.getLogger(__name__) _PYTHON_VERSION_FILES = [ Path("mlflow", "version.py"), ] _PYPROJECT_TOML_FILES = [ Path("pyproject.toml"), Path("pyproject.release.toml"), Path("libs/skinny/pyproject.toml"), Path("libs/tracing/pyproject.toml"), ] _JAVA_VERSION_FILES = list(Path("mlflow", "java").rglob("*.java")) _JAVA_POM_XML_FILES = list(Path("mlflow", "java").rglob("*.xml")) _TS_VERSION_FILES = [ Path( "mlflow", "server", "js", "src", "common", "constants.tsx", ), Path("docs", "src", "constants.ts"), ] _R_VERSION_FILES = [Path("mlflow", "R", "mlflow", "DESCRIPTION")] _HELM_CHART_FILES = [Path("charts", "Chart.yaml")] def get_current_py_version() -> str: text = Path("mlflow", "version.py").read_text() match = re.search(r'VERSION = "(.+)"', text) if match is None: raise ValueError("Could not find VERSION in mlflow/version.py") return match.group(1) def get_java_py_version_pattern(version: str) -> str: version_without_suffix = replace_dev_or_rc_suffix_with(version, "") return rf"{re.escape(version_without_suffix)}(-SNAPSHOT)?" def get_java_new_py_version(new_py_version: str) -> str: return replace_dev_or_rc_suffix_with(new_py_version, "-SNAPSHOT") def replace_dev_or_rc_suffix_with(version: str, repl: str) -> str: parsed = Version(version) base_version = parsed.base_version return base_version + repl if parsed.is_prerelease else version def replace_occurrences(files: list[Path], pattern: str | re.Pattern[str], repl: str) -> None: if not isinstance(pattern, re.Pattern): pattern = re.compile(pattern) for f in files: old_text = f.read_text() if not pattern.search(old_text): continue new_text = pattern.sub(repl, old_text) f.write_text(new_text) def replace_python(old_version: str, new_py_version: str, paths: list[Path]) -> None: replace_occurrences( files=paths, pattern=re.escape(old_version), repl=new_py_version, ) def replace_pyproject_toml(new_py_version: str, paths: list[Path]) -> None: replace_occurrences( files=paths, pattern=re.compile(r'^version\s+=\s+".+"$', re.MULTILINE), repl=f'version = "{new_py_version}"', ) # Update mlflow-skinny and mlflow-tracing versions to match the new mlflow version. replace_occurrences( files=paths, pattern=re.compile(r"^\s*\"mlflow-skinny==.+\",$", re.MULTILINE), repl=f' "mlflow-skinny=={new_py_version}",', ) replace_occurrences( files=paths, pattern=re.compile(r"^\s*\"mlflow-tracing==.+\",$", re.MULTILINE), repl=f' "mlflow-tracing=={new_py_version}",', ) def replace_ts(old_version: str, new_py_version: str, paths: list[Path]) -> None: replace_occurrences( files=paths, pattern=re.escape(old_version), repl=new_py_version, ) def replace_java(old_version: str, new_py_version: str, paths: list[Path]) -> None: old_py_version_pattern = get_java_py_version_pattern(old_version) dev_suffix_replaced = get_java_new_py_version(new_py_version) replace_occurrences( files=paths, pattern=old_py_version_pattern, repl=dev_suffix_replaced, ) # Note: the pom.xml files define versions of dependencies as # well. this causes issues when the mlflow version matches the # version of a dependency. to work around, we make sure to # match only the correct keys def replace_java_pom_xml(old_version: str, new_py_version: str, paths: list[Path]) -> None: old_py_version_pattern = get_java_py_version_pattern(old_version) dev_suffix_replaced = get_java_new_py_version(new_py_version) mlflow_version_tag_pattern = r"" mlflow_spark_pattern = r"mlflow-spark_2\.1[23]\s+" mlflow_parent_pattern = r"mlflow-parent\s+" # combine the three tags together to form the regex mlflow_replace_pattern = ( rf"({mlflow_version_tag_pattern}|{mlflow_spark_pattern}|{mlflow_parent_pattern})" + f"{old_py_version_pattern}" + r"(|)" ) # group 1: everything before the version # group 2: optional -SNAPSHOT # group 3: everything after the version replace_str = f"\\g<1>{dev_suffix_replaced}\\g<3>" replace_occurrences( files=paths, pattern=mlflow_replace_pattern, repl=replace_str, ) def replace_helm_chart(new_py_version: str, paths: list[Path]) -> None: replace_occurrences( files=paths, pattern=re.compile(r'^appVersion:\s+".+"$', re.MULTILINE), repl=f'appVersion: "{replace_dev_or_rc_suffix_with(new_py_version, "")}"', ) def replace_r(old_py_version: str, new_py_version: str, paths: list[Path]) -> None: current_py_version_without_suffix = replace_dev_or_rc_suffix_with(old_py_version, "") replace_occurrences( files=paths, pattern=f"Version: {re.escape(current_py_version_without_suffix)}", repl=f"Version: {replace_dev_or_rc_suffix_with(new_py_version, '')}", ) def update_versions(new_py_version: str, helm_app_version: str | None = None) -> None: """ `new_py_version` is either: - a release version (e.g. "2.1.0") - a RC version (e.g. "2.1.0rc0") - a dev version (e.g. "2.1.0.dev0") `helm_app_version` is the version written to the Helm chart's `appVersion`. The Helm chart's `appVersion` is used as the default Docker image tag, so it must point to a published release. When omitted, the Helm chart is left untouched. """ old_py_version = get_current_py_version() replace_python(old_py_version, new_py_version, _PYTHON_VERSION_FILES) replace_pyproject_toml(new_py_version, _PYPROJECT_TOML_FILES) replace_ts(old_py_version, new_py_version, _TS_VERSION_FILES) replace_java(old_py_version, new_py_version, _JAVA_VERSION_FILES) replace_java_pom_xml(old_py_version, new_py_version, _JAVA_POM_XML_FILES) replace_r(old_py_version, new_py_version, _R_VERSION_FILES) if helm_app_version is not None: replace_helm_chart(helm_app_version, _HELM_CHART_FILES) def validate_new_version(value: str) -> str: new = Version(value) current = Version(get_current_py_version()) # this could be the case if we just promoted an RC to a release if new < current: _logger.warning( f"New version {new} is not greater than or equal to current version {current}. " "If the previous version was an RC, this is expected. If not, please make sure the " "specified new version is correct." ) # exit with 0 to avoid failing the CI job exit(0) return value def pre_release(new_version: str) -> None: """ Update MLflow package versions BEFORE release. Usage: python dev/update_mlflow_versions.py pre-release --new-version 1.29.0 """ validate_new_version(new_version) update_versions(new_py_version=new_version) def post_release(new_version: str) -> None: """ Update MLflow package versions AFTER release. Usage: python dev/update_mlflow_versions.py post-release --new-version 1.29.0 """ validate_new_version(new_version) current_version = Version(get_current_py_version()) msg = ( "It appears you ran this command on a release branch because the current version " f"({current_version}) is not a dev version. Please re-run this command on the master " "branch." ) assert current_version.is_devrelease, msg new_version_parsed = Version(new_version) # Increment the patch version and append ".dev0" new_py_version = ( f"{new_version_parsed.major}.{new_version_parsed.minor}.{new_version_parsed.micro + 1}.dev0" ) # Helm `appVersion` must point to the released version, not the next dev # version, since it's used as the default Docker image tag. update_versions(new_py_version=new_py_version, helm_app_version=new_version) if __name__ == "__main__": parser = argparse.ArgumentParser(description="Update MLflow package versions") subparsers = parser.add_subparsers(dest="command", help="Command to run", required=True) # pre-release subcommand pre_parser = subparsers.add_parser( "pre-release", help="Update MLflow package versions BEFORE release", ) pre_parser.add_argument("--new-version", required=True, help="New version to release") # post-release subcommand post_parser = subparsers.add_parser( "post-release", help="Update MLflow package versions AFTER release", ) post_parser.add_argument("--new-version", required=True, help="New version that was released") args = parser.parse_args() if args.command == "pre-release": pre_release(args.new_version) elif args.command == "post-release": post_release(args.new_version)