Repository: fal-ai/isolate (language: python)

Issue:
feat: Implement a Pyenv backend (FEA-566)
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #52
* #51
* __->__ #50




Discussion / hints:

Fix patch:
diff --git a/pyproject.toml b/pyproject.toml
index 3d29519..502c298 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,6 +22,7 @@ server = []
 "conda" = "isolate.backends.conda:CondaEnvironment"
 "local" = "isolate.backends.local:LocalPythonEnvironment"
 "isolate-server" = "isolate.backends.remote:IsolateServer"
+"pyenv" = "isolate.backends.pyenv:PyenvEnvironment"
 
 [build-system]
 requires = ["poetry-core>=1.1.0"]
diff --git a/src/isolate/backends/pyenv.py b/src/isolate/backends/pyenv.py
new file mode 100644
index 0000000..82ffc22
--- /dev/null
+++ b/src/isolate/backends/pyenv.py
@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import functools
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, ClassVar, Dict, List, Optional
+
+from isolate.backends import BaseEnvironment, EnvironmentCreationError
+from isolate.backends.common import active_python, logged_io, sha256_digest_of
+from isolate.backends.settings import DEFAULT_SETTINGS, IsolateSettings
+from isolate.connections import PythonIPC
+
+_PYENV_EXECUTABLE_NAME = "pyenv"
+_PYENV_EXECUTABLE_PATH = os.environ.get("ISOLATE_PYENV_EXECUTABLE")
+
+
+@dataclass
+class PyenvEnvironment(BaseEnvironment[Path]):
+    BACKEND_NAME: ClassVar[str] = "pyenv"
+
+    python_version: str
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        settings: IsolateSettings = DEFAULT_SETTINGS,
+    ) -> BaseEnvironment:
+        environment = cls(**config)
+        environment.apply_settings(settings)
+        return environment
+
+    @property
+    def key(self) -> str:
+        return os.path.join("versions", self.python_version)
+
+    def create(self) -> Path:
+        pyenv = _get_pyenv_executable()
+        env_path = self.settings.cache_dir_for(self)
+        with self.settings.cache_lock_for(env_path):
+            # PyEnv installs* the Python versions under $root/versions/$version, where we
+            # use versions/$version as the key and $root as the base directory (for pyenv).
+            #
+            # [0]: https://github.com/pyenv/pyenv#locating-pyenv-provided-python-installations
+            pyenv_root = env_path.parent.parent
+            prefix = self._try_get_prefix(pyenv, pyenv_root)
+            if prefix is None:
+                self._install_python(pyenv, pyenv_root)
+                prefix = self._try_get_prefix(pyenv, pyenv_root)
+                if not prefix:
+                    raise EnvironmentCreationError(
+                        f"Python {self.python_version} must have been installed by now."
+                    )
+
+            assert prefix is not None
+            return prefix
+
+    def _try_get_prefix(self, pyenv: Path, root_path: Path) -> Optional[Path]:
+        try:
+            prefix = subprocess.check_output(
+                [pyenv, "prefix", self.python_version],
+                env={**os.environ, "PYENV_ROOT": str(root_path)},
+                text=True,
+                stderr=subprocess.PIPE,
+            )
+        except subprocess.CalledProcessError as exc:
+            if "not installed" in exc.stderr:
+                return None
+            raise EnvironmentCreationError(
+                f"Failed to get the prefix for Python {self.python_version}.\n"
+                f"{exc.stdout}\n{exc.stderr}"
+            )
+
+        return Path(prefix.strip())
+
+    def _install_python(self, pyenv: Path, root_path: Path) -> None:
+        with logged_io(self.log) as (stdout, stderr):
+            try:
+                subprocess.check_call(
+                    [pyenv, "install", "--skip-existing", self.python_version],
+                    env={**os.environ, "PYENV_ROOT": str(root_path)},
+                    stdout=stdout,
+                    stderr=stderr,
+                )
+            except subprocess.CalledProcessError as exc:
+                raise EnvironmentCreationError(
+                    f"Failed to install Python {self.python_version} via pyenv.\n"
+                )
+
+    def destroy(self, connection_key: Path) -> None:
+        pyenv = _get_pyenv_executable()
+        with self.settings.cache_lock_for(connection_key):
+            # It might be destroyed already (when we are awaiting
+            # for the lock to be released).
+            if not connection_key.exists():
+                return None
+
+            pyenv_root = connection_key.parent.parent
+            with logged_io(self.log) as (stdout, stderr):
+                subprocess.check_call(
+                    [pyenv, "uninstall", "-f", connection_key.name],
+                    env={**os.environ, "PYENV_ROOT": str(pyenv_root)},
+                    stdout=stdout,
+                    stderr=stderr,
+                )
+
+    def exists(self) -> bool:
+        pyenv = _get_pyenv_executable()
+        cache_dir = self.settings.cache_dir_for(self)
+        with self.settings.cache_lock_for(cache_dir):
+            pyenv_root = cache_dir.parent.parent
+            prefix = self._try_get_prefix(pyenv, pyenv_root)
+            return prefix is not None
+
+    def open_connection(self, connection_key: Path) -> PythonIPC:
+        return PythonIPC(self, connection_key)
+
+
+@functools.lru_cache(1)
+def _get_pyenv_executable() -> Path:
+    if _PYENV_EXECUTABLE_PATH:
+        if not os.path.exists(_PYENV_EXECUTABLE_PATH):
+            raise EnvironmentCreationError(
+                f"Path to pyenv executable not found! ISOLATE_PYENV_EXECUTABLE variable: {_PYENV_EXECUTABLE_PATH!r}"
+            )
+        return Path(_PYENV_EXECUTABLE_PATH)
+
+    pyenv_path = shutil.which(_PYENV_EXECUTABLE_NAME)
+    if pyenv_path is None:
+        raise FileNotFoundError(
+            "Could not find the pyenv executable. If pyenv is not already installed in your system, please"
+            "install it first. If it is not in your PATH, then point ISOLATE_PYENV_COMMAND to the absolute path of the"
+            "pyenv executable."
+        )
+    return Path(pyenv_path)


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 26b5ba2..1718cc6 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -36,10 +36,18 @@ jobs:
         # PyJokes is available on conda-forge
         conda-channels: anaconda, conda-forge
 
+    - uses: actions/checkout@v3
+      with:
+        repository: pyenv/pyenv
+        ref: v2.3.6
+        path: pyenv
+
     - name: Install dependencies
       run: |
         python -m pip install -r dev-requirements.txt
         python -m pip install -e ".[server]"
 
     - name: Test
-      run: python -m pytest
+      run: |
+        export ISOLATE_PYENV_EXECUTABLE=pyenv/bin/pyenv
+        python -m pytest
diff --git a/tests/test_backends.py b/tests/test_backends.py
index 9f9a7cd..868057a 100644
--- a/tests/test_backends.py
+++ b/tests/test_backends.py
@@ -14,6 +14,7 @@
 from isolate.backends.common import sha256_digest_of
 from isolate.backends.conda import CondaEnvironment, _get_conda_executable
 from isolate.backends.local import LocalPythonEnvironment
+from isolate.backends.pyenv import PyenvEnvironment, _get_pyenv_executable
 from isolate.backends.remote import IsolateServer
 from isolate.backends.settings import IsolateSettings
 from isolate.backends.virtualenv import VirtualPythonEnvironment
@@ -295,7 +296,7 @@ def test_custom_python_version(self, tmp_path):
 
 
 # Since conda is an external dependency, we'll skip tests using it
-# if it is not installed. v1 vbn
+# if it is not installed.
 try:
     _get_conda_executable()
 except FileNotFoundError:
@@ -589,3 +590,40 @@ def test_isolate_server_multiple_envs(isolate_server):
 def test_wrong_options(kind, config):
     with pytest.raises(TypeError):
         isolate.prepare_environment(kind, **config)
+
+
+# Since pyenv is an external dependency, we'll skip tests using it
+# if it is not installed.
+try:
+    _get_pyenv_executable()
+except FileNotFoundError:
+    IS_PYENV_AVAILABLE = False
+else:
+    IS_PYENV_AVAILABLE = True
+
+
+@pytest.mark.skipif(not IS_PYENV_AVAILABLE, reason="Pyenv is not available")
+@pytest.mark.parametrize("python_version", ["3.7", "3.10", "3.9.15"])
+def test_pyenv_environment(python_version, tmp_path):
+    different_python = PyenvEnvironment(python_version)
+    test_settings = IsolateSettings(Path(tmp_path))
+    different_python.apply_settings(test_settings)
+
+    assert not different_python.exists()
+
+    connection_key = different_python.create()
+    assert different_python.exists()
+
+    with different_python.open_connection(connection_key) as connection:
+        assert connection.run(partial(eval, "__import__('sys').version")).startswith(
+            python_version
+        )
+
+    # Do a cached run.
+    with different_python.open_connection(connection_key) as connection:
+        assert connection.run(partial(eval, "__import__('sys').version")).startswith(
+            python_version
+        )
+
+    different_python.destroy(connection_key)
+    assert not different_python.exists()