Share a venv, if possible (#155)

This also includes a few small fixes.

Author: ericsnowcurrently

pyperformance/_pip.py

@@ -36,26 +36,32 @@ def run_pip(cmd, *args, **kwargs):
 
 def is_pip_installed(python, *, env=None):
     """Return True if pip is installed on the given Python executable."""
-    ec, _, _ = run_pip('--version', env=env, capture=True, verbose=False)
+    ec, _, _ = run_pip(
+        '--version',
+        python=python,
+        env=env,
+        capture=True,
+        verbose=False,
+    )
     return ec == 0
 
 
 def install_pip(python=sys.executable, *,
                 info=None,
                 downloaddir=None,
                 env=None,
+                upgrade=True,
                 **kwargs
                 ):
     """Install pip on the given Python executable."""
     if not python:
         python = getattr(info, 'executable', None) or sys.executable
 
     # python -m ensurepip
-    res = _utils.run_python(
-        '-m', 'ensurepip', '--verbose',
-        python=python,
-        **kwargs
-    )
+    args = ['-m', 'ensurepip', '-v']  # --verbose
+    if upgrade:
+        args.append('-U')  # --upgrade
+    res = _utils.run_python(*args, python=python, **kwargs)
     ec, _, _ = res
     if ec == 0 and is_pip_installed(python, env=env):
         return res
@@ -111,16 +117,20 @@ def upgrade_pip(python=sys.executable, *,
 
     if installer:
         # Upgrade installer dependencies (setuptools, ...)
-        reqs = [
-            f'setuptools>={OLD_SETUPTOOLS}',
-            # install wheel so pip can cache binary wheel packages locally,
-            # and install prebuilt wheel packages from PyPI.
-            'wheel',
-        ]
-        res = install_requirements(*reqs, python=python, upgrade=True, **kwargs)
+        res = ensure_installer(python, upgrade=True, **kwargs)
     return res
 
 
+def ensure_installer(python=sys.executable, **kwargs):
+    reqs = [
+        f'setuptools>={OLD_SETUPTOOLS}',
+        # install wheel so pip can cache binary wheel packages locally,
+        # and install prebuilt wheel packages from PyPI.
+        'wheel',
+    ]
+    return install_requirements(*reqs, python=python, **kwargs)
+
+
 def install_requirements(reqs, *extra,
                          upgrade=True,
                          **kwargs

pyperformance/_venv.py

@@ -197,18 +197,41 @@ def base(self):
             self._base = _pythoninfo.get_info(base_exe)
             return self._base
 
-    def ensure_pip(self, downloaddir=None):
+    def ensure_pip(self, downloaddir=None, *, installer=True, upgrade=True):
+        if not upgrade and _pip.is_pip_installed(self.python, env=self._env):
+            return
         ec, _, _ = _pip.install_pip(
             self.python,
             info=self.info,
             downloaddir=downloaddir or self.root,
             env=self._env,
+            upgrade=upgrade,
         )
         if ec != 0:
             raise VenvPipInstallFailedError(root, ec)
         elif not _pip.is_pip_installed(self.python, env=self._env):
             raise VenvPipInstallFailedError(root, 0, "pip doesn't work")
 
+        if installer:
+            # Upgrade installer dependencies (setuptools, ...)
+            ec, _, _ = _pip.ensure_installer(
+                self.python,
+                env=self._env,
+                upgrade=True,
+            )
+            if ec != 0:
+                raise RequirementsInstallationFailedError('wheel')
+
+    def upgrade_pip(self, *, installer=True):
+        ec, _, _ = _pip.upgrade_pip(
+            self.python,
+            info=self.info,
+            env=self._env,
+            installer=installer,
+        )
+        if ec != 0:
+            raise RequirementsInstallationFailedError('pip')
+
     def ensure_reqs(self, *reqs, upgrade=True):
         print("Installing requirements into the virtual environment %s" % self.root)
         ec, _, _ = _pip.install_requirements(

pyperformance/run.py

@@ -60,10 +60,16 @@ def run_benchmarks(should_run, python, options):
     info = _pythoninfo.get_info(python)
     runid = get_run_id(info)
 
+    unique = getattr(options, 'unique_venvs', False)
+    if not unique:
+        common = VenvForBenchmarks.ensure(
+            _venv.get_venv_root(runid.name, python=info),
+            info,
+            inherit_environ=options.inherit_environ,
+        )
+
     benchmarks = {}
     venvs = set()
-    if sys.prefix != sys.base_prefix:
-        venvs.add(sys.prefix)
     for i, bench in enumerate(to_run):
         bench_runid = runid._replace(bench=bench)
         assert bench_runid.name, (bench, bench_runid)
@@ -73,12 +79,20 @@ def run_benchmarks(should_run, python, options):
         print('='*50)
         print(f'({i+1:>2}/{len(to_run)}) creating venv for benchmark ({bench.name})')
         print()
-        alreadyseen = venv_root in venvs
+        if not unique:
+            print('(trying common venv first)')
+            # Try the common venv first.
+            try:
+                common.ensure_reqs(bench)
+            except _venv.RequirementsInstallationFailedError:
+                print('(falling back to unique venv)')
+            else:
+                benchmarks[bench] = (common, bench_runid)
+                continue
         venv = VenvForBenchmarks.ensure(
             venv_root,
             info,
             inherit_environ=options.inherit_environ,
-            refresh=not alreadyseen,
         )
         try:
             # XXX Do not override when there is a requirements collision.
@@ -89,6 +103,7 @@ def run_benchmarks(should_run, python, options):
             venv = None
         venvs.add(venv_root)
         benchmarks[bench] = (venv, bench_runid)
+    print()
 
     suite = None
     run_count = str(len(to_run))

pyperformance/venv.py

@@ -120,7 +120,7 @@ class VenvForBenchmarks(_venv.VirtualEnvironment):
     @classmethod
     def create(cls, root=None, python=None, *,
                inherit_environ=None,
-               install=True,
+               upgrade=False,
                ):
         env = _get_envvars(inherit_environ)
         try:
@@ -131,7 +131,7 @@ def create(cls, root=None, python=None, *,
         self.inherit_environ = inherit_environ
 
         try:
-            self.ensure_pip()
+            self.ensure_pip(upgrade=upgrade)
         except _venv.VenvPipInstallFailedError as exc:
             print(f'ERROR: {exc}')
             _utils.safe_rmtree(self.root)
@@ -140,92 +140,64 @@ def create(cls, root=None, python=None, *,
             _utils.safe_rmtree(self.root)
             raise
 
-        try:
-            self.prepare(install)
-        except BaseException:
-            print()
-            _utils.safe_rmtree(venv.root)
-            raise
+        # Display the pip version
+        _pip.run_pip('--version', python=self.python, env=self._env)
 
         return self
 
     @classmethod
     def ensure(cls, root, python=None, *,
-               refresh=True,
-               install=True,
+               upgrade=False,
                **kwargs
                ):
         if _venv.venv_exists(root):
             self = super().ensure(root)
-            if refresh:
-                self.prepare(install)
+            if upgrade:
+                self.upgrade_pip()
+            else:
+                self.ensure_pip(upgrade=False)
             return self
         else:
-            return cls.create(root, python, install=install, **kwargs)
+            return cls.create(
+                root,
+                python,
+                upgrade=upgrade,
+                **kwargs
+            )
 
     def __init__(self, root, *, base=None, inherit_environ=None):
         super().__init__(root, base=base)
         self.inherit_environ = inherit_environ or None
-        self._prepared = False
 
     @property
     def _env(self):
         # Restrict the env we use.
         return _get_envvars(self.inherit_environ)
 
-    def prepare(self, install=True):
-        print("Installing the virtual environment %s" % self.root)
-        if self._prepared or (self._prepared is None and not install):
-            print('(already installed)')
-            return
-
-        if not self._prepared:
-            # parse requirements
+    def install_pyperformance(self):
+        print("installing pyperformance in the venv at %s" % self.root)
+        # Install pyperformance inside the virtual environment.
+        if pyperformance.is_dev():
             basereqs = Requirements.from_file(REQUIREMENTS_FILE, ['psutil'])
+            self.ensure_reqs(basereqs)
 
-            # Upgrade pip
-            ec, _, _ = _pip.upgrade_pip(
-                self.python,
-                info=self.info,
+            root_dir = os.path.dirname(pyperformance.PKG_ROOT)
+            ec, _, _ = _pip.install_editable(
+                root_dir,
+                python=self.info,
                 env=self._env,
-                installer=True,
             )
-            if ec != 0:
-                sys.exit(ec)
-
-        # XXX not for benchmark venvs
-        if install:
-            # install pyperformance inside the virtual environment
-            # XXX This isn't right...
-            if pyperformance.is_installed():
-                root_dir = os.path.dirname(pyperformance.PKG_ROOT)
-                ec, _, _ = _pip.install_editable(
-                    root_dir,
-                    python=self.info,
-                    env=self._env,
-                )
-            else:
-                version = pyperformance.__version__
-                ec, _, _ = _pip.install_requirements(
-                    f'pyperformance=={version}',
-                    python=self.info,
-                    env=self._env,
-                )
-            if ec != 0:
-                sys.exit(ec)
-            self._prepared = True
         else:
-            self._prepared = None
-
-        # Display the pip version
-        _pip.run_pip('--version', python=self.python, env=self._env)
-
-        # Dump the package list and their versions: pip freeze
-        _pip.run_pip('freeze', python=self.python, env=self._env)
+            version = pyperformance.__version__
+            ec, _, _ = _pip.install_requirements(
+                f'pyperformance=={version}',
+                python=self.info,
+                env=self._env,
+            )
+        if ec != 0:
+            sys.exit(ec)
 
     def ensure_reqs(self, requirements=None, *, exitonerror=False):
-        print("Installing requirements into the virtual environment %s" % self.root)
-
         # parse requirements
         bench = None
         if requirements is None:
@@ -235,20 +207,17 @@ def ensure_reqs(self, requirements=None, *, exitonerror=False):
             requirements = Requirements.from_benchmarks([bench])
 
         # Every benchmark must depend on pyperf.
-        if requirements and bench is not None:
-            if not requirements.get('pyperf'):
-                basereqs = Requirements.from_file(REQUIREMENTS_FILE, ['psutil'])
-                pyperf_req = basereqs.get('pyperf')
-                if not pyperf_req:
-                    raise NotImplementedError
-                requirements.specs.append(pyperf_req)
-                # XXX what about psutil?
+        if bench is not None and not requirements.get('pyperf'):
+            basereqs = Requirements.from_file(REQUIREMENTS_FILE, ['psutil'])
+            pyperf_req = basereqs.get('pyperf')
+            if not pyperf_req:
+                raise NotImplementedError
+            requirements.specs.append(pyperf_req)
+            # XXX what about psutil?
 
         if not requirements:
             print('(nothing to install)')
         else:
-            self.prepare(install=bench is None)
-
             # install requirements
             try:
                 super().ensure_reqs(
@@ -287,6 +256,12 @@ def cmd_venv(options, benchmarks=None):
             info = None
     exists = _venv.venv_exists(root)
 
+    def install(venv):
+        try:
+            venv.install_pyperformance()
+        except _venv.RequirementsInstallationFailedError:
+            sys.exit(1)
+
     action = options.venv_action
     if action == 'create':
         requirements = Requirements.from_benchmarks(benchmarks)
@@ -297,6 +272,8 @@ def cmd_venv(options, benchmarks=None):
             info,
             inherit_environ=options.inherit_environ,
         )
+        venv.ensure_pip()
+        install(venv)
         venv.ensure_reqs(requirements, exitonerror=True)
         if not exists:
             print("The virtual environment %s has been created" % root)
@@ -313,6 +290,8 @@ def cmd_venv(options, benchmarks=None):
                     info,
                     inherit_environ=options.inherit_environ,
                 )
+                venv.ensure_pip()
+                install(venv)
                 venv.ensure_reqs(requirements, exitonerror=True)
             else:
                 print("The virtual environment %s already exists" % root)
@@ -324,6 +303,8 @@ def cmd_venv(options, benchmarks=None):
                     info,
                     inherit_environ=options.inherit_environ,
                 )
+                venv.ensure_pip()
+                install(venv)
                 venv.ensure_reqs(requirements, exitonerror=True)
                 print("The virtual environment %s has been recreated" % root)
         else:
@@ -332,6 +313,8 @@ def cmd_venv(options, benchmarks=None):
                 info,
                 inherit_environ=options.inherit_environ,
             )
+            venv.ensure_pip()
+            install(venv)
             venv.ensure_reqs(requirements, exitonerror=True)
             print("The virtual environment %s has been created" % root)