New tool-agnostic ASV environment management

benchmarks/asv.conf.json

@@ -3,18 +3,25 @@
     "project": "scitools-iris",
     "project_url": "https://github.com/SciTools/iris",
     "repo": "..",
-    "environment_type": "nox-conda",
+    "environment_type": "conda-delegated",
     "show_commit_url": "https://github.com/scitools/iris/commit/",
 
     "benchmark_dir": "./benchmarks",
     "env_dir": ".asv/env",
     "results_dir": ".asv/results",
     "html_dir": ".asv/html",
-    "plugins": [".nox_asv_plugin"],
-    // The commit to checkout to first run Nox to set up the environment.
-    "nox_setup_commit": "HEAD",
-    // The path of the noxfile's location relative to the project root.
-    "noxfile_rel_path": "noxfile.py",
-    // The ``--session`` arg to be used with ``--install-only`` to prep an environment.
-    "nox_session_name": "tests"
+    "plugins": [".asv_delegated_conda"],
+
+    // The command(s) that create/update an environment correctly for the
+    //  checked-out commit.
+    // Interpreted the same as build_command, with following exceptions:
+    //  * No build-time environment variables.
+    //  * Is run in the same environment as the ASV install itself.
+    "delegated_env_commands": [
+        "sed -i 's/_PY_VERSIONS_ALL/_PY_VERSION_LATEST/g' noxfile.py",
+        "nox --envdir={conf_dir}/.asv/env/nox01 --session=tests --install-only --no-error-on-external-run --verbose"
+    ],
+    // The parent directory of the above environment.
+    //  The most recently modified environment in the directory will be used.
+    "delegated_env_parent": "{conf_dir}/.asv/env/nox01"
 }

benchmarks/asv_delegated_conda.py

@@ -0,0 +1,208 @@
+# Copyright Iris contributors
+#
+# This file is part of Iris and is released under the LGPL license.
+# See COPYING and COPYING.LESSER in the root of the repository for full
+# licensing details.
+"""
+ASV plug-in providing an alternative :class:`asv.plugins.conda.Conda`
+subclass that manages the Conda environment via custom user scripts.
+
+"""
+
+from os import environ
+from os.path import getmtime
+from pathlib import Path
+from shutil import copy2, copytree, rmtree
+from tempfile import TemporaryDirectory
+
+from asv import util as asv_util
+from asv.config import Config
+from asv.console import log
+from asv.plugins.conda import Conda
+from asv.repo import Repo
+
+
+class CondaDelegated(Conda):
+    """
+    Manage a Conda environment using custom user scripts, run at each commit.
+
+    Ignores user input variations - ``matrix`` / ``pythons`` /
+    ``conda_environment_file``, since environment is being managed outside ASV.
+
+    Original environment creation behaviour is inherited, but upon checking out
+    a commit the custom script(s) are run and the original environment is
+    replaced with a symlink to the custom environment. This arrangement is then
+    re-used in subsequent runs.
+
+    """
+
+    tool_name = "conda-delegated"
+
+    def __init__(
+        self,
+        conf: Config,
+        python: str,
+        requirements: dict,
+        tagged_env_vars: dict,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        conf : Config instance
+
+        python : str
+            Version of Python.  Must be of the form "MAJOR.MINOR".
+
+        requirements : dict
+            Dictionary mapping a PyPI package name to a version
+            identifier string.
+
+        tagged_env_vars : dict
+            Environment variables, tagged for build vs. non-build
+
+        """
+        ignored = ["`python`"]
+        if requirements:
+            ignored.append("`requirements`")
+        if tagged_env_vars:
+            ignored.append("`tagged_env_vars`")
+        if conf.conda_environment_file:
+            ignored.append("`conda_environment_file`")
+        message = (
+            f"Ignoring ASV setting(s): {', '.join(ignored)}. Benchmark "
+            "environment management is delegated to third party script(s)."
+        )
+        log.warning(message)
+        requirements = {}
+        tagged_env_vars = {}
+        conf.conda_environment_file = None
+
+        super().__init__(conf, python, requirements, tagged_env_vars)
+        self._update_info()
+
+        self._env_commands = self._interpolate_commands(
+            conf.delegated_env_commands
+        )
+        # Again using _interpolate_commands to get env parent path - allows use
+        #  of the same ASV env variables.
+        env_parent_interpolated = self._interpolate_commands(
+            conf.delegated_env_parent
+        )
+        # Returns list of tuples, we just want the first.
+        env_parent_first = env_parent_interpolated[0]
+        # The 'command' is the first item in the returned tuple.
+        env_parent_string = " ".join(env_parent_first[0])
+        self._delegated_env_parent = Path(env_parent_string).resolve()
+
+    @property
+    def name(self):
+        """Get a name to uniquely identify this environment."""
+        return asv_util.sanitize_filename(self.tool_name)
+
+    def _update_info(self) -> None:
+        """Make sure class properties reflect the actual environment being used."""
+        # Follow symlink if it has been created.
+        actual_path = Path(self._path).resolve()
+        self._path = str(actual_path)
+
+        # Get custom environment's Python version if it exists yet.
+        try:
+            get_version = (
+                "from sys import version_info; "
+                "print(f'{version_info.major}.{version_info.minor}')"
+            )
+            actual_python = self.run(["-c", get_version])
+            self._python = actual_python
+        except OSError:
+            pass
+
+    def _prep_env(self) -> None:
+        """Run the custom environment script(s) and switch to using that environment."""
+        message = f"Running delegated environment management for: {self.name}"
+        log.info(message)
+        env_path = Path(self._path)
+
+        def copy_asv_files(src_parent: Path, dst_parent: Path) -> None:
+            """For copying between self._path and a temporary cache."""
+            asv_files = list(src_parent.glob("asv*"))
+            # build_root_path.name usually == "project" .
+            asv_files += [src_parent / Path(self._build_root).name]
+            for src_path in asv_files:
+                dst_path = dst_parent / src_path.name
+                if not dst_path.exists():
+                    # Only caching in case the environment has been rebuilt.
+                    #  If the dst_path already exists: rebuilding hasn't
+                    #  happened. Also a non-issue when copying in the reverse
+                    #  direction because the cache dir is temporary.
+                    if src_path.is_dir():
+                        func = copytree
+                    else:
+                        func = copy2
+                    func(src_path, dst_path)
+
+        with TemporaryDirectory(prefix="delegated_asv_cache_") as asv_cache:
+            asv_cache_path = Path(asv_cache)
+            # Cache all of ASV's files as delegated command may remove and
+            #  re-build the environment.
+            copy_asv_files(env_path.resolve(), asv_cache_path)
+
+            # Adapt the build_dir to the cache location.
+            build_root_path = Path(self._build_root)
+            build_dir_original = build_root_path / self._repo_subdir
+            build_dir_subpath = build_dir_original.relative_to(
+                build_root_path.parent
+            )
+            build_dir = asv_cache_path / build_dir_subpath
+
+            # Run the script(s) for delegated environment creation/updating.
+            # (An adaptation of self._interpolate_and_run_commands).
+            for command, env, return_codes, cwd in self._env_commands:
+                local_envs = dict(environ)
+                local_envs.update(env)
+                if cwd is None:
+                    cwd = str(build_dir)
+                _ = asv_util.check_output(
+                    command,
+                    timeout=self._install_timeout,
+                    cwd=cwd,
+                    env=local_envs,
+                    valid_return_codes=return_codes,
+                )
+
+            # Replace the env that ASV created with a symlink to the env
+            #  created/updated by the custom script.
+            delegated_env_path = sorted(
+                self._delegated_env_parent.glob("*"),
+                key=getmtime,
+                reverse=True,
+            )[0]
+            if env_path.resolve() != delegated_env_path:
+                try:
+                    env_path.unlink(missing_ok=True)
+                except IsADirectoryError:
+                    rmtree(env_path)
+                env_path.symlink_to(
+                    delegated_env_path, target_is_directory=True
+                )
+
+            # Check that environment exists.
+            try:
+                env_path.resolve(strict=True)
+            except FileNotFoundError:
+                message = f"Path does not resolve to environment: {env_path}"
+                log.error(message)
+                raise RuntimeError(message)
+
+            # Restore ASV's files from the cache (if necessary).
+            copy_asv_files(asv_cache_path, env_path.resolve())
+
+            # Record new environment information in properties.
+            self._update_info()
+
+    def checkout_project(self, repo: Repo, commit_hash: str) -> None:
+        """Check out the working tree of the project at given commit hash."""
+        super().checkout_project(repo, commit_hash)
+        self._prep_env()
+        log.info(
+            f"Environment {self.name} updated to spec at {commit_hash[:8]}"
+        )