update_environments_and_lock_files.py

"""Script to update CI environment files and associated lock files.

To run it you need to be in the root folder of the scikit-learn repo:
python build_tools/update_environments_and_lock_files.py

Two scenarios where this script can be useful:
- make sure that the latest versions of all the dependencies are used in the CI.
  We can run this script regularly and open a PR with the changes to the lock
  files. This workflow will eventually be automated with a bot in the future.
- bump minimum dependencies in sklearn/_min_dependencies.py. Running this
  script will update both the CI environment files and associated lock files.
  You can then open a PR with the changes.
- pin some packages to an older version by adding them to the
  default_package_constraints variable. This is useful when regressions are
  introduced in our dependencies, this has happened for example with pytest 7
  and coverage 6.3.

Environments are conda environment.yml or pip requirements.txt. Lock files are
conda-lock lock files or pip-compile requirements.txt.

pip requirements.txt are used when we install some dependencies (e.g. numpy and
scipy) with apt-get and the rest of the dependencies (e.g. pytest and joblib)
with pip.

To run this script you need:
- conda-lock. The version should match the one used in the CI in
  sklearn/_min_dependencies.py
- pip-tools

"""

import re
import subprocess
import sys
from pathlib import Path
import shlex
import json
import logging
from importlib.metadata import version

import click

from conda_lock import __version__ as conda_lock_version
from jinja2 import Environment
from packaging.version import Version

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
logger.addHandler(handler)


common_dependencies_without_coverage = [
    "python",
    "numpy",
    "blas",
    "scipy",
    "cython",
    "joblib",
    "threadpoolctl",
    "matplotlib",
    "pandas",
    "pyamg",
    "pytest",
    "pytest-xdist",
    "pillow",
]

common_dependencies = common_dependencies_without_coverage + [
    "pytest-cov",
    "coverage",
]

docstring_test_dependencies = ["sphinx", "numpydoc"]

default_package_constraints = {
    # XXX: pin pytest-xdist to workaround:
    # https://github.com/pytest-dev/pytest-xdist/issues/840
    "pytest-xdist": "2.5.0",
}


def remove_from(alist, to_remove):
    return [each for each in alist if each not in to_remove]


conda_build_metadata_list = [
    {
        "build_name": "pylatest_conda_forge_mkl_linux-64",
        "folder": "build_tools/azure",
        "platform": "linux-64",
        "channel": "conda-forge",
        "conda_dependencies": common_dependencies + [
            "ccache",
            "pytorch",
            "pytorch-cpu",
            "array-api-compat",
        ],
        "package_constraints": {
            "blas": "[build=mkl]",
            "pytorch": "1.13",
        },
    },
    {
        "build_name": "pylatest_conda_forge_mkl_osx-64",
        "folder": "build_tools/azure",
        "platform": "osx-64",
        "channel": "conda-forge",
        "conda_dependencies": common_dependencies + [
            "ccache",
            "compilers",
            "llvm-openmp",
        ],
        "package_constraints": {
            "blas": "[build=mkl]",
        },
    },
    {
        "build_name": "pylatest_conda_mkl_no_openmp",
        "folder": "build_tools/azure",
        "platform": "osx-64",
        "channel": "defaults",
        # TODO work-around to get cython>=0.29.33 via PyPi until it is in conda defaults
        # See: https://github.com/ContinuumIO/anaconda-issues/issues/13120
        "conda_dependencies": remove_from(common_dependencies, ["cython"]) + ["ccache"],
        # TODO work-around to get cython>=0.29.33 via PyPi until it is in conda defaults
        # See: https://github.com/ContinuumIO/anaconda-issues/issues/13120
        "pip_dependencies": ["cython"],
        "package_constraints": {
            "blas": "[build=mkl]",
            # 2022-06-09 currently mamba install 1.23 and scipy 1.7 which
            # should be compatible but actually are not. This pin can be
            # removed when scipy 1.8 is available in conda defaults channel.
            # For more details, see
            # https://github.com/scikit-learn/scikit-learn/pull/24363#issuecomment-1236927660
            # and https://github.com/scipy/scipy/issues/16964
            "numpy": "1.22",
            # XXX: coverage is temporary pinned to 6.2 because 6.3 is not
            # fork-safe and 6.4 is not available yet (July 2022) in conda
            # defaults channel. For more details, see:
            # https://github.com/nedbat/coveragepy/issues/1310
            "coverage": "6.2",
        },
    },
    {
        "build_name": "pylatest_conda_forge_mkl_no_coverage",
        "folder": "build_tools/azure",
        "platform": "linux-64",
        "channel": "conda-forge",
        "conda_dependencies": common_dependencies_without_coverage + ["ccache"],
        "package_constraints": {
            "blas": "[build=mkl]",
        },
    },
    {
        "build_name": "py38_conda_defaults_openblas",
        "folder": "build_tools/azure",
        "platform": "linux-64",
        "channel": "defaults",
        # TODO work-around to get cython>=0.29.33 via PyPi until it is in conda defaults
        # See: https://github.com/ContinuumIO/anaconda-issues/issues/13120
        "conda_dependencies": remove_from(common_dependencies, ["cython"]) + ["ccache"],
        # TODO work-around to get cython>=0.29.33 via PyPi until it is in conda defaults
        # See: https://github.com/ContinuumIO/anaconda-issues/issues/13120
        "pip_dependencies": ["cython"],
        "package_constraints": {
            "python": "3.8",
            "blas": "[build=openblas]",
            "numpy": "min",
            "scipy": "min",
            "matplotlib": "min",
            "threadpoolctl": "2.2.0",
            # XXX: coverage is temporary pinned to 6.2 because 6.3 is not
            # fork-safe and 6.4 is not available yet (July 2022) in conda
            # defaults channel. For more details, see:
            # https://github.com/nedbat/coveragepy/issues/1310
            "coverage": "6.2",
        },
    },
    {
        "build_name": "py38_conda_forge_openblas_ubuntu_2204",
        "folder": "build_tools/azure",
        "platform": "linux-64",
        "channel": "conda-forge",
        "conda_dependencies": common_dependencies_without_coverage + ["ccache"],
        "package_constraints": {"python": "3.8", "blas": "[build=openblas]"},
    },
    {
        "build_name": "pylatest_pip_openblas_pandas",
        "folder": "build_tools/azure",
        "platform": "linux-64",
        "channel": "defaults",
        # sphinx in conda_dependencies as a temporary work-around for
        # https://github.com/conda-incubator/conda-lock/issues/309
        "conda_dependencies": ["python", "ccache", "sphinx"],
        "pip_dependencies": (
            remove_from(common_dependencies, ["python", "blas"])
            + remove_from(docstring_test_dependencies, ["sphinx"])
            + ["lightgbm", "scikit-image"]
        ),
        "package_constraints": {
            "python": "3.9",
            # TODO: remove the following pin when we have a solution since
            # pinning a pandas version partially defeats the purpose of the
            # `pylastest` CI configurations which are meant to test
            # scikit-learn against the most recently released versions of
            # its dependencies.
            # pandas 2.0 depends on tzdata as well as python installed via conda
            # We currently have a mix of conda/pip and run into the following issue:
            # https://github.com/conda/conda-lock/issues/179
            "pandas": "1.5.3",
        },
    },
    {
        "build_name": "pylatest_pip_scipy_dev",
        "folder": "build_tools/azure",
        "platform": "linux-64",
        "channel": "defaults",
        # sphinx in conda_dependencies as a temporary work-around for
        # https://github.com/conda-incubator/conda-lock/issues/309
        "conda_dependencies": ["python", "ccache", "sphinx"],
        "pip_dependencies": (
            remove_from(
                common_dependencies,
                [
                    "python",
                    "blas",
                    "matplotlib",
                    "pyamg",
                    # all the dependencies below have a development version
                    # installed in the CI, so they can be removed from the
                    # environment.yml
                    "numpy",
                    "scipy",
                    "pandas",
                    "cython",
                    "joblib",
                    "pillow",
                ],
            )
            + ["pooch"]
            + remove_from(docstring_test_dependencies, ["sphinx"])
            # python-dateutil is a dependency of pandas and pandas is removed from
            # the environment.yml. Adding python-dateutil so it is pinned
            + ["python-dateutil"]
        ),
    },
    {
        "build_name": "pypy3",
        "folder": "build_tools/azure",
        "platform": "linux-64",
        "channel": "conda-forge",
        "conda_dependencies": (
            ["pypy", "python"]
            + remove_from(
                common_dependencies_without_coverage, ["python", "pandas", "pillow"]
            )
            + ["ccache"]
        ),
        "package_constraints": {
            "blas": "[build=openblas]",
            "python": "3.9",
        },
    },
    {
        "build_name": "py38_conda_forge_mkl",
        "folder": "build_tools/azure",
        "platform": "win-64",
        "channel": "conda-forge",
        "conda_dependencies": remove_from(common_dependencies, ["pandas", "pyamg"]) + [
            "wheel",
            "pip",
        ],
        "package_constraints": {
            "python": "3.8",
            "blas": "[build=mkl]",
        },
    },
    {
        "build_name": "doc_min_dependencies",
        "folder": "build_tools/circle",
        "platform": "linux-64",
        "channel": "conda-forge",
        "conda_dependencies": common_dependencies_without_coverage + [
            "scikit-image",
            "seaborn",
            "memory_profiler",
            "compilers",
            "sphinx",
            "sphinx-gallery",
            "numpydoc",
            "sphinx-prompt",
            "plotly",
            "pooch",
        ],
        "pip_dependencies": ["sphinxext-opengraph"],
        "package_constraints": {
            "python": "3.8",
            "numpy": "min",
            "scipy": "min",
            "matplotlib": "min",
            "cython": "min",
            "scikit-image": "min",
            "sphinx": "min",
            "pandas": "min",
            "sphinx-gallery": "min",
            "numpydoc": "min",
            "sphinx-prompt": "min",
            "sphinxext-opengraph": "min",
            "plotly": "min",
        },
    },
    {
        "build_name": "doc",
        "folder": "build_tools/circle",
        "platform": "linux-64",
        "channel": "conda-forge",
        "conda_dependencies": common_dependencies_without_coverage + [
            "scikit-image",
            "seaborn",
            "memory_profiler",
            "compilers",
            "sphinx",
            "sphinx-gallery",
            "numpydoc",
            "sphinx-prompt",
            "plotly",
            "pooch",
            "sphinxext-opengraph",
        ],
        "pip_dependencies": ["jupyterlite-sphinx", "jupyterlite-pyodide-kernel"],
        "package_constraints": {
            "python": "3.9",
            # XXX: sphinx > 6.0 does not correctly generate searchindex.js
            "sphinx": "6.0.0",
        },
    },
    {
        "build_name": "py39_conda_forge",
        "folder": "build_tools/cirrus",
        "platform": "linux-aarch64",
        "channel": "conda-forge",
        "conda_dependencies": remove_from(
            common_dependencies_without_coverage, ["pandas", "pyamg"]
        ) + ["pip", "ccache"],
        "package_constraints": {
            "python": "3.9",
        },
    },
]


pip_build_metadata_list = [
    {
        "build_name": "debian_atlas_32bit",
        "folder": "build_tools/azure",
        "pip_dependencies": [
            "cython",
            "joblib",
            "threadpoolctl",
            "pytest",
            "pytest-cov",
        ],
        "package_constraints": {
            "joblib": "min",
            "threadpoolctl": "2.2.0",
            "pytest": "min",
            "pytest-cov": "min",
            # no pytest-xdist because it causes issue on 32bit
        },
        # same Python version as in debian-32 build
        "python_version": "3.9.2",
    },
    {
        "build_name": "ubuntu_atlas",
        "folder": "build_tools/azure",
        "pip_dependencies": [
            "cython",
            "joblib",
            "threadpoolctl",
            "pytest",
            "pytest-xdist",
        ],
        "package_constraints": {"joblib": "min", "threadpoolctl": "min"},
        "python_version": "3.10.4",
    },
]


def execute_command(command_list):
    proc = subprocess.Popen(
        command_list, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )

    out, err = proc.communicate()
    out, err = out.decode(), err.decode()

    if proc.returncode != 0:
        command_str = " ".join(command_list)
        raise RuntimeError(
            "Command exited with non-zero exit code.\n"
            "Exit code: {}\n"
            "Command:\n{}\n"
            "stdout:\n{}\n"
            "stderr:\n{}\n".format(proc.returncode, command_str, out, err)
        )
    return out


def get_package_with_constraint(package_name, build_metadata, uses_pip=False):
    build_package_constraints = build_metadata.get("package_constraints")
    if build_package_constraints is None:
        constraint = None
    else:
        constraint = build_package_constraints.get(package_name)

    constraint = constraint or default_package_constraints.get(package_name)

    if constraint is None:
        return package_name

    comment = ""
    if constraint == "min":
        constraint = execute_command(
            [sys.executable, "sklearn/_min_dependencies.py", package_name]
        ).strip()
        comment = "  # min"

    if re.match(r"\d[.\d]*", constraint):
        equality = "==" if uses_pip else "="
        constraint = equality + constraint

    return f"{package_name}{constraint}{comment}"


environment = Environment(trim_blocks=True, lstrip_blocks=True)
environment.filters["get_package_with_constraint"] = get_package_with_constraint


def get_conda_environment_content(build_metadata):
    template = environment.from_string("""
# DO NOT EDIT: this file is generated from the specification found in the
# following script to centralize the configuration for CI builds:
# build_tools/update_environments_and_lock_files.py
channels:
  - {{ build_metadata['channel'] }}
dependencies:
  {% for conda_dep in build_metadata['conda_dependencies'] %}
  - {{ conda_dep | get_package_with_constraint(build_metadata) }}
  {% endfor %}
  {% if build_metadata['pip_dependencies'] %}
  - pip
  - pip:
  {% for pip_dep in build_metadata.get('pip_dependencies', []) %}
    - {{ pip_dep | get_package_with_constraint(build_metadata, uses_pip=True) }}
  {% endfor %}
  {% endif %}""".strip())
    return template.render(build_metadata=build_metadata)


def write_conda_environment(build_metadata):
    content = get_conda_environment_content(build_metadata)
    build_name = build_metadata["build_name"]
    folder_path = Path(build_metadata["folder"])
    output_path = folder_path / f"{build_name}_environment.yml"
    output_path.write_text(content)


def write_all_conda_environments(build_metadata_list):
    for build_metadata in build_metadata_list:
        write_conda_environment(build_metadata)


def conda_lock(environment_path, lock_file_path, platform):
    command = (
        f"conda-lock lock --mamba --kind explicit --platform {platform} "
        f"--file {environment_path} --filename-template {lock_file_path}"
    )

    logger.debug("conda-lock command: %s", command)
    execute_command(shlex.split(command))


def create_conda_lock_file(build_metadata):
    build_name = build_metadata["build_name"]
    folder_path = Path(build_metadata["folder"])
    environment_path = folder_path / f"{build_name}_environment.yml"
    platform = build_metadata["platform"]
    lock_file_basename = build_name
    if not lock_file_basename.endswith(platform):
        lock_file_basename = f"{lock_file_basename}_{platform}"

    lock_file_path = folder_path / f"{lock_file_basename}_conda.lock"
    conda_lock(environment_path, lock_file_path, platform)


def write_all_conda_lock_files(build_metadata_list):
    for build_metadata in build_metadata_list:
        logger.info(build_metadata["build_name"])
        create_conda_lock_file(build_metadata)


def get_pip_requirements_content(build_metadata):
    template = environment.from_string("""
# DO NOT EDIT: this file is generated from the specification found in the
# following script to centralize the configuration for CI builds:
# build_tools/update_environments_and_lock_files.py
{% for pip_dep in build_metadata['pip_dependencies'] %}
{{ pip_dep | get_package_with_constraint(build_metadata, uses_pip=True) }}
{% endfor %}""".strip())
    return template.render(build_metadata=build_metadata)


def write_pip_requirements(build_metadata):
    build_name = build_metadata["build_name"]
    content = get_pip_requirements_content(build_metadata)
    folder_path = Path(build_metadata["folder"])
    output_path = folder_path / f"{build_name}_requirements.txt"
    output_path.write_text(content)


def write_all_pip_requirements(build_metadata_list):
    for build_metadata in build_metadata_list:
        logger.info(build_metadata["build_name"])
        write_pip_requirements(build_metadata)


def pip_compile(pip_compile_path, requirements_path, lock_file_path):
    command = f"{pip_compile_path} --upgrade {requirements_path} -o {lock_file_path}"

    logger.debug("pip-compile command: %s", command)
    execute_command(shlex.split(command))


def write_pip_lock_file(build_metadata):
    build_name = build_metadata["build_name"]
    python_version = build_metadata["python_version"]
    environment_name = f"pip-tools-python{python_version}"
    # To make sure that the Python used to create the pip lock file is the same
    # as the one used during the CI build where the lock file is used, we first
    # create a conda environment with the correct Python version and
    # pip-compile and run pip-compile in this environment

    command = (
        "conda create -c conda-forge -n"
        f" pip-tools-python{python_version} python={python_version} pip-tools -y"
    )
    execute_command(shlex.split(command))

    json_output = execute_command(shlex.split("conda info --json"))
    conda_info = json.loads(json_output)
    environment_folder = [
        each for each in conda_info["envs"] if each.endswith(environment_name)
    ][0]
    environment_path = Path(environment_folder)
    pip_compile_path = environment_path / "bin" / "pip-compile"

    folder_path = Path(build_metadata["folder"])
    requirement_path = folder_path / f"{build_name}_requirements.txt"
    lock_file_path = folder_path / f"{build_name}_lock.txt"
    pip_compile(pip_compile_path, requirement_path, lock_file_path)


def write_all_pip_lock_files(build_metadata_list):
    for build_metadata in build_metadata_list:
        write_pip_lock_file(build_metadata)


def remove_unnecessary_package_from_lock_file(build_metadata_list, build_name, package):
    # find the build metadata from the list
    build_metadata = None
    for metadata in build_metadata_list:
        if metadata["build_name"] == build_name:
            build_metadata = metadata
            break
    if build_metadata is None:
        logger.warning(
            f"Could not find {build_name} in build_metadata_list. This may be expected"
            " if --select-build is used"
        )
        return
    folder_path = Path(build_metadata["folder"])
    platform = build_metadata["platform"]
    lock_file_basename = build_name
    if not lock_file_basename.endswith(platform):
        lock_file_basename = f"{lock_file_basename}_{platform}"

    lock_file_path = folder_path / f"{lock_file_basename}_conda.lock"
    with open(lock_file_path, "r") as f_read_only:
        lock_file_content = f_read_only.readlines()
        with open(lock_file_path, "w") as f_write:
            for line in lock_file_content:
                if package in line:
                    # simple heuristic to remove the package
                    # it will not work if package that have partial string in common
                    logger.debug(f"Removing {package} from {build_name} lock file")
                    logger.debug(f"Removed line: {line}")
                    continue
                f_write.write(line)


def check_conda_lock_version():
    # Check that the installed conda-lock version is consistent with _min_dependencies.
    expected_conda_lock_version = execute_command(
        [sys.executable, "sklearn/_min_dependencies.py", "conda-lock"]
    ).strip()

    installed_conda_lock_version = version("conda-lock")
    if installed_conda_lock_version != expected_conda_lock_version:
        raise RuntimeError(
            f"Expected conda-lock version: {expected_conda_lock_version}, got:"
            f" {installed_conda_lock_version}"
        )


@click.command()
@click.option(
    "--select-build",
    default="",
    help="Regex to restrict the builds we want to update environment and lock files",
)
def main(select_build):
    check_conda_lock_version()
    filtered_conda_build_metadata_list = [
        each
        for each in conda_build_metadata_list
        if re.search(select_build, each["build_name"])
    ]
    logger.info("Writing conda environments")
    write_all_conda_environments(filtered_conda_build_metadata_list)
    logger.info("Writing conda lock files")
    write_all_conda_lock_files(filtered_conda_build_metadata_list)
    if Version(conda_lock_version) < Version("1.5"):
        # In conda-lock < 1.5, `os_name` was not considered when creating the
        # dependency trees:
        # https://github.com/conda/conda-lock/issues/293
        # Here, we need to manually remove the package `pywinpty` from the `doc`
        # build which is only needed on Windows required by `jupyter-server`.
        remove_unnecessary_package_from_lock_file(
            filtered_conda_build_metadata_list, "doc", "pywinpty"
        )

    filtered_pip_build_metadata_list = [
        each
        for each in pip_build_metadata_list
        if re.search(select_build, each["build_name"])
    ]
    logger.info("Writing pip requirements")
    write_all_pip_requirements(filtered_pip_build_metadata_list)
    logger.info("Writing pip lock files")
    write_all_pip_lock_files(filtered_pip_build_metadata_list)


if __name__ == "__main__":
    main()