feat: conda env for bentos in bentostore (#3396)

* feat: conda env for bentos in bentostore

* give a bento store

* fix info messages

* loading ephimeral env for path_to_bento

* fix typing

* log error if return_code!=0

* test for env_manager decorator utils

* get python version

* attrs -> attr

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

* address review comments

* added and Environment class to handle environment specific operations

* remove pipefail

* polish debug mode

* documentation

* added to serve-grpc

Update src/bentoml_cli/env_manager.py

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

Update docs/source/guides/envmanager.rst

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

Update docs/source/guides/envmanager.rst

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

Update docs/source/guides/envmanager.rst

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

Update docs/source/guides/envmanager.rst

Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

changes

* fix commits from review

* use bento internally in environment

* rename EnvironmentFactory to EnvManager

---------

Co-authored-by: Jithin James <jjmachan@pop-os.localdomain>
Co-authored-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>

docs/source/guides/envmanager.rst

@@ -0,0 +1,50 @@
+===================
+Environment Manager
+===================
+
+:bdg-info:`Note:` This feature is currently only supported on UNIX/MacOS.
+
+Environment manager is a utility that helps create an isolated environment to
+run the BentoML CLI. Dependencies are pulled from your defined
+``bentofile.yaml`` and the environment is built upon request. This means by
+passing ``--env`` to supported CLI commands (such as :ref:`bentoml serve
+<reference/cli:serve>`), such commands will then be run in an sandbox
+environment that mimics the behaviour during production.
+
+.. code-block:: bash
+
+   » bentoml serve --env conda iris_classifier:latest
+
+This creates and isolated conda environment from the dependencies in the bento
+and runs ``bentoml serve`` from that environment.
+
+.. note:: The current implementation will try to install the given dependencies
+   before running the CLI command. Therefore, the environment startup will be a
+   blocking call.
+
+
+BentoML CLI Commands that support Environment Manager
+    - :ref:`serve <reference/cli:serve>`
+    - :ref:`serve-grpc <reference/cli:serve-grpc>`
+
+Supported Environments
+    - conda
+
+
+Caching strategies
+==================
+
+Currently, there are two types of environments that are supported by the
+environment manager:
+
+1. Persistent environment: If the given target is a Bento, then the created
+   environment will be stored locally to ``$BENTOML_HOME/env``. Such an
+   environment will then be cached and later used by subsequent invocations.
+
+2. Ephemeral environment: In cases where the given target is not a Bento (import
+   path to ``bentoml.Service``, project directory containing a valid
+   ``bentofile.yaml``), the environment will be created and cleanup up on
+   demand.
+
+.. note::
+   You can run ``rm -rf $BENTOML_HOME/env`` to clear the cache.

docs/source/guides/index.rst

@@ -80,6 +80,7 @@ into this part of the documentation.
     client
     server
     configuration
+    envmanager
     graph
     monitoring
     logging

src/bentoml/_internal/configuration/containers.py

@@ -29,6 +29,8 @@
 from ..utils.unflatten import unflatten
 
 if TYPE_CHECKING:
+    from fs.base import FS
+
     from .. import external_typing as ext
     from ..models import ModelStore
     from ..utils.analytics import ServeInfo
@@ -160,7 +162,6 @@ def to_dict(self) -> providers.ConfigDictType:
 
 @dataclass
 class _BentoMLContainerClass:
-
     config = providers.Configuration()
 
     @providers.SingletonFactory
@@ -175,8 +176,9 @@ def bentoml_home() -> str:
         )
         bentos = os.path.join(home, "bentos")
         models = os.path.join(home, "models")
+        envs = os.path.join(home, "envs")
 
-        validate_or_create_dir(home, bentos, models)
+        validate_or_create_dir(home, bentos, models, envs)
         return home
 
     @providers.SingletonFactory
@@ -189,6 +191,18 @@ def bento_store_dir(bentoml_home: str = Provide[bentoml_home]):
     def model_store_dir(bentoml_home: str = Provide[bentoml_home]):
         return os.path.join(bentoml_home, "models")
 
+    @providers.SingletonFactory
+    @staticmethod
+    def env_store_dir(bentoml_home: str = Provide[bentoml_home]):
+        return os.path.join(bentoml_home, "envs")
+
+    @providers.SingletonFactory
+    @staticmethod
+    def env_store(bentoml_home: str = Provide[bentoml_home]) -> FS:
+        import fs
+
+        return fs.open_fs(os.path.join(bentoml_home, "envs"))
+
     @providers.SingletonFactory
     @staticmethod
     def bento_store(base_dir: str = Provide[bento_store_dir]):

src/bentoml/_internal/env_manager/__init__.py

@@ -0,0 +1,3 @@
+from .manager import EnvManager
+
+__all__ = ["EnvManager"]

src/bentoml/_internal/env_manager/envs.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import os
+import typing as t
+import logging
+import subprocess
+from abc import ABC
+from abc import abstractmethod
+from shutil import which
+from tempfile import NamedTemporaryFile
+
+import attr
+
+from ...exceptions import BentoMLException
+from ..bento.bento import Bento
+from ..configuration import get_debug_mode
+
+logger = logging.getLogger(__name__)
+
+if t.TYPE_CHECKING:
+    from fs.base import FS
+
+
+@attr.define
+class Environment(ABC):
+    name: str
+    env_fs: FS
+    # path to bento's /env dir
+    bento: Bento
+    env_exe: str = attr.field(init=False)
+
+    def __attrs_post_init__(self):
+        self.env_exe = self.get_executable()
+        if self.env_fs.exists(self.name):
+            return self.env_fs.getsyspath(self.name)
+        else:
+            self.create()
+
+    @abstractmethod
+    def get_executable(self) -> str:
+        """
+        Returns executable path responsible for running this environment.
+
+        Make sure that said executable is available in PATH.
+        """
+        ...
+
+    @abstractmethod
+    def create(self):
+        """
+        Create the environment with the files from bento.
+        """
+        ...
+
+    @abstractmethod
+    def run(self, commands: list[str]):
+        """
+        run the commands in an activated environment.
+        """
+        ...
+
+    @staticmethod
+    def run_script_subprocess(
+        script_file_path: t.Union[str, os.PathLike[str]],
+        capture_output: bool,
+        debug_mode: bool,
+    ):
+        shell_path = which("bash")
+        if shell_path is None:
+            raise BentoMLException("Unable to locate a valid shell")
+
+        safer_bash_args: list[str] = []
+
+        # but only work in bash
+        if debug_mode:
+            safer_bash_args = ["-euxo", "pipefail"]
+        result = subprocess.run(
+            [shell_path, *safer_bash_args, script_file_path],
+            capture_output=capture_output,
+        )
+        if result.returncode != 0:
+            if result.stdout:
+                logger.debug(result.stdout.decode())
+            if result.stderr:
+                logger.error(result.stderr.decode())
+            raise BentoMLException(
+                "Subprocess call returned non-zero value. Reffer logs for more details"
+            )
+
+
+class Conda(Environment):
+    def get_executable(self) -> str:
+        conda_exe = os.environ.get("CONDA_EXE")
+        if conda_exe is None:
+            raise BentoMLException(
+                "Conda|Miniconda executable not found! Make sure any one is installed and environment is activated."
+            )
+        return conda_exe
+
+    def create(self):
+        # install conda deps
+        from ..bento.build_config import CONDA_ENV_YAML_FILE_NAME
+
+        # create a env under $BENTOML_HOME/env
+        # setup conda with bento's environment.yml file and python/install.sh file
+        conda_env_path = self.env_fs.getsyspath(self.name)
+        python_version: str
+        with open(self.bento.path_of("/env/python/version.txt"), "r") as pyver_file:
+            py_version = pyver_file.read().split(".")[:2]
+            python_version = ".".join(py_version)
+        conda_environment_file = self.bento.path_of(
+            f"/env/conda/{CONDA_ENV_YAML_FILE_NAME}"
+        )
+        create_script = f"""\
+eval "$({self.get_executable()} shell.posix hook)"
+
+conda create -p {conda_env_path} python={python_version} --yes
+
+if [ -f {conda_environment_file} ]; then
+    conda config --set pip_interop_enabled True
+    conda env update -p {conda_env_path} --file {conda_environment_file}
+fi
+
+conda activate {conda_env_path}
+bash -euxo pipefail {self.bento.path_of("/env/python/install.sh")}
+"""
+        with NamedTemporaryFile(mode="w", delete=False) as script_file:
+            script_file.write(create_script)
+
+        logger.info("Creating Conda env and installing dependencies...")
+        self.run_script_subprocess(
+            script_file.name,
+            capture_output=not get_debug_mode(),
+            debug_mode=get_debug_mode(),
+        )
+
+    def run(self, commands: list[str]):
+        """
+        Run commands in the activated environment.
+        """
+        conda_env_path = self.env_fs.getsyspath(self.name)
+        create_script = f"""\
+eval "$({self.env_exe} shell.posix hook)"
+conda activate {conda_env_path}
+{" ".join(commands)}
+"""
+        with NamedTemporaryFile(mode="w", delete=False) as script_file:
+            script_file.write(create_script)
+        self.run_script_subprocess(
+            script_file.name, capture_output=False, debug_mode=get_debug_mode()
+        )

src/bentoml/_internal/env_manager/manager.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import typing as t
+import logging
+
+import fs
+from simple_di import inject
+from simple_di import Provide
+
+from .envs import Conda
+from ..bento.bento import Bento
+from ..bento.bento import BentoInfo
+from ..configuration.containers import BentoMLContainer
+
+if t.TYPE_CHECKING:
+    from fs.base import FS
+
+    from .envs import Environment
+
+logger = logging.getLogger(__name__)
+
+
+class EnvManager:
+    environment: Environment
+
+    @inject
+    def __init__(
+        self,
+        env_type: t.Literal["conda"],
+        bento: Bento,
+        is_ephemeral: bool = True,
+        env_name: str | None = None,
+        env_store: FS = Provide[BentoMLContainer.env_store],
+    ):
+        if not is_ephemeral:
+            assert env_name is not None, "persistent environments need a valid name."
+        if not env_store.exists(env_type):
+            env_store.makedir(env_type)
+        env_fs = fs.open_fs("temp://") if is_ephemeral else env_store.opendir(env_type)
+
+        if env_type == "conda":
+            self.environment = Conda(name=env_name, env_fs=env_fs, bento=bento)
+        else:
+            raise NotImplementedError(f"'{env_type}' is not supported.")
+
+    @classmethod
+    def from_bento(
+        cls,
+        env_type: t.Literal["conda"],
+        bento: Bento,
+        is_ephemeral: bool,
+    ) -> EnvManager:
+        env_name: str
+        if is_ephemeral:
+            env_name = "ephemeral_env"
+        else:
+            env_name = str(bento.tag).replace(":", "_")
+        return cls(
+            env_type=env_type,
+            env_name=env_name,
+            is_ephemeral=is_ephemeral,
+            bento=bento,
+        )
+
+    @classmethod
+    def from_bentofile(
+        cls, env_type: str, bento_info: BentoInfo, is_ephemeral: str
+    ) -> EnvManager:
+        raise NotImplementedError

src/bentoml/bentos.py

@@ -430,6 +430,7 @@ def serve(
     server_type: str = "http",
     reload: bool = False,
     production: bool = False,
+    env: t.Literal["conda"] | None = None,
     host: str | None = None,
     port: int | None = None,
     working_dir: str | None = None,
@@ -503,6 +504,8 @@ def serve(
         args.append("--production")
     if reload:
         args.append("--reload")
+    if env:
+        args.extend(["--env", env])
 
     if api_workers is not None:
         args.extend(["--api-workers", str(api_workers)])