Skip to content
Permalink
Browse files

Add pip_dependencies as parameter for env decorator

Add pip_dependencies as parameter.

```
from bentoml import env

@env(pip_dependencies=['numpy', 'pandas']
```
  • Loading branch information...
yubozhao committed Jul 30, 2019
1 parent b3dc158 commit df20db24b5d6c4a63b510f623adbe9815aea2f9e
Showing with 87 additions and 13 deletions.
  1. +2 −1 bentoml/service.py
  2. +21 −12 bentoml/service_env.py
  3. +64 −0 tests/test_service_env.py
@@ -262,7 +262,8 @@ def env_decorator(**kwargs):
setup_sh (str): User defined shell script to run before running BentoService.
It could be local file path or the shell script content.
requirements_text (str): User defined requirement text to install before
running BentoService. It could be local file path or requirements' content
running BentoService.
pip_dependencies (str or list(str)): User defined python modules to install.
conda_channels (list(str)): User defined conda channels
conda_dependencies (list(str)): Defined dependencies to be installed with
conda environment.
@@ -114,7 +114,7 @@ class BentoServiceEnv(object):
def __init__(self):
self._setup_sh = None
self._conda_env = CondaEnv()
self._requirements_txt = None
self._pip_dependencies = []

def get_conda_env_name(self):
return self._conda_env.get_name()
@@ -140,7 +140,7 @@ def add_conda_pip_dependencies(self, pip_dependencies):
def add_handler_dependencies(self, handler_dependencies):
if not isinstance(handler_dependencies, list):
handler_dependencies = [handler_dependencies]
self._conda_env.add_pip_dependencies(handler_dependencies)
self._pip_dependencies += handler_dependencies

def set_setup_sh(self, setup_sh_path_or_content):
setup_sh_file = Path(setup_sh_path_or_content)
@@ -151,22 +151,28 @@ def set_setup_sh(self, setup_sh_path_or_content):
else:
self._setup_sh = setup_sh_path_or_content.encode("utf-8")

def set_requirements_txt(self, requirements_txt_path_or_content):
requirements_txt_file = Path(requirements_txt_path_or_content)
def add_pip_dependencies(self, pip_dependencies):
if not isinstance(pip_dependencies, list):
pip_dependencies = [pip_dependencies]
self._pip_dependencies += pip_dependencies

if requirements_txt_file.is_file():
with requirements_txt_file.open("rb") as f:
self._requirements_txt = f.read()
else:
self._requirements_txt = requirements_txt_path_or_content.encode("utf-8")
def set_requirements_txt(self, requirements_txt_path):
requirements_txt_file = Path(requirements_txt_path)

with requirements_txt_file.open("rb") as f:
content = f.read()
module_list = content.decode("utf-8").split("\n")
self._pip_dependencies += module_list

def save(self, path):
conda_yml_file = os.path.join(path, "environment.yml")
self._conda_env.write_to_yaml_file(conda_yml_file)

requirements_txt_file = os.path.join(path, "requirements.txt")

with open(requirements_txt_file, "wb") as f:
f.write(self._requirements_txt or b"")
pip_content = "\n".join(self._pip_dependencies).encode("utf-8")
f.write(pip_content)

if self._setup_sh:
setup_sh_file = os.path.join(path, "setup.sh")
@@ -183,6 +189,9 @@ def from_dict(cls, env_dict):
if "requirements_txt" in env_dict:
env.set_requirements_txt(env_dict["requirements_txt"])

if "pip_dependencies" in env_dict:
env.add_pip_dependencies(env_dict["pip_dependencies"])

if "conda_channels" in env_dict:
env.add_conda_channels(env_dict["conda_channels"])

@@ -200,8 +209,8 @@ def to_dict(self):
if self._setup_sh:
env_dict["setup_sh"] = self._setup_sh

if self._requirements_txt:
env_dict["requirements_txt"] = self._requirements_txt
if self._pip_dependencies:
env_dict["pip_dependencies"] = self._pip_dependencies

env_dict["conda_env"] = self._conda_env._conda_env

@@ -0,0 +1,64 @@
import os

import bentoml


def test_requirement_txt_env(tmpdir):
req_txt_file = tmpdir.join("requirements.txt")
model = ''
with open(str(req_txt_file), 'wb') as f:
f.write(b"numpy\npandas\ntorch")

@bentoml.artifacts([bentoml.artifact.PickleArtifact('model')])
@bentoml.env(requirements_txt=str(req_txt_file))
class ServiceWithFile(bentoml.BentoService):
@bentoml.api(bentoml.handlers.DataframeHandler)
def predict(self, df):
return df

service_with_file = ServiceWithFile.pack(model=model)
assert len(service_with_file.env._pip_dependencies) == 3


def test_pip_dependencies_env(tmpdir):
model = ''

@bentoml.artifacts([bentoml.artifact.PickleArtifact('model')])
@bentoml.env(pip_dependencies="numpy")
class ServiceWithString(bentoml.BentoService):
@bentoml.api(bentoml.handlers.DataframeHandler)
def predict(self, df):
return df

@bentoml.artifacts([bentoml.artifact.PickleArtifact('model')])
@bentoml.env(pip_dependencies=['numpy', 'pandas', 'torch'])
class ServiceWithList(bentoml.BentoService):
@bentoml.api(bentoml.handlers.DataframeHandler)
def predict(self, df):
return df

service_with_string = ServiceWithString.pack(model=model)
assert 'numpy' in service_with_string.env._pip_dependencies

service_with_list = ServiceWithList.pack(model=model)
assert len(service_with_list.env._pip_dependencies) == 3


def test_pip_dependencies_with_archive(tmpdir):
model = ''

@bentoml.artifacts([bentoml.artifact.PickleArtifact('model')])
@bentoml.env(pip_dependencies=['numpy', 'pandas', 'torch'])
class ServiceWithList(bentoml.BentoService):
@bentoml.api(bentoml.handlers.DataframeHandler)
def predict(self, df):
return df

service_with_list = ServiceWithList.pack(model=model)
saved_path = service_with_list.save(tmpdir)

requirements_txt_path = os.path.join(saved_path, 'requirements.txt')
with open(requirements_txt_path, 'rb') as f:
saved_requirements = f.read()
module_list = saved_requirements.decode('utf-8').split('\n')
assert len(module_list) == 3

0 comments on commit df20db2

Please sign in to comment.
You can’t perform that action at this time.