feat(venv): support pip options (#1052)

* Support pip options when creating venv. Support extra env_vars. Reformat Python files

* ci: auto fixes from pre-commit.ci

For more information, see https://pre-commit.ci

* Update local.py

Signed-off-by: Rick Zhou <rickzhoucmu@gmail.com>

* ci: auto fixes from pre-commit.ci

For more information, see https://pre-commit.ci

---------

Signed-off-by: Rick Zhou <rickzhoucmu@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/openllm/analytic.py

@@ -104,7 +104,7 @@ def wrapped(ctx: click.Context, *args, **kwargs):
                                 cmd_name=command_name,
                                 duration_in_ms=duration_in_ns / 1e6,
                                 error_type=type(e).__name__,
-                                return_code=2 if isinstance(e, KeyboardInterrupt) else 1,
+                                return_code=(2 if isinstance(e, KeyboardInterrupt) else 1),
                             )
                         )
                         raise

src/openllm/common.py

@@ -167,6 +167,10 @@ def version(self) -> str:
     def labels(self) -> dict[str, str]:
         return self.bento_yaml['labels']
 
+    @property
+    def envs(self) -> list[dict[str, str]]:
+        return self.bento_yaml['envs']
+
     @functools.cached_property
     def bento_yaml(self) -> dict:
         import yaml
@@ -227,7 +231,8 @@ def tolist(self):
 
 class VenvSpec(SimpleNamespace):
     python_version: str
-    python_packages: dict[str, str]
+    python_packages: list[str]
+    options: list[str] = []
     name_prefix = ''
 
     def __hash__(self):

src/openllm/local.py

@@ -7,6 +7,18 @@
 from openllm.venv import ensure_venv
 
 
+def prep_env_vars(bento: BentoInfo):
+    import os
+
+    env_vars = bento.envs
+    for env_var in env_vars:
+        if 'value' not in env_var:
+            continue
+        key = env_var['name']
+        value = env_var['value']
+        os.environ[key] = value
+
+
 def _get_serve_cmd(bento: BentoInfo, port: int = 3000):
     cmd = ['bentoml', 'serve', bento.bentoml_tag]
     if port != 3000:
@@ -16,6 +28,7 @@ def _get_serve_cmd(bento: BentoInfo, port: int = 3000):
 
 
 def serve(bento: BentoInfo, port: int = 3000):
+    prep_env_vars(bento)
     venv = ensure_venv(bento)
     cmd, env, cwd = _get_serve_cmd(bento, port=port)
     output(f'Access the Chat UI at http://localhost:{port}/chat (or with you IP)')
@@ -91,4 +104,5 @@ async def _run_model(bento: BentoInfo, port: int = 3000, timeout: int = 600):
 
 
 def run(bento: BentoInfo, port: int = 3000, timeout: int = 600):
+    prep_env_vars(bento)
     asyncio.run(_run_model(bento, port=port, timeout=timeout))

src/openllm/repo.py

@@ -8,53 +8,43 @@
 import typer
 
 from openllm.analytic import OpenLLMTyper
-from openllm.common import (
-    INTERACTIVE,
-    REPO_DIR,
-    VERBOSE_LEVEL,
-    RepoInfo,
-    load_config,
-    output,
-    save_config,
-)
+from openllm.common import INTERACTIVE, REPO_DIR, VERBOSE_LEVEL, RepoInfo, load_config, output, save_config
 
 UPDATE_INTERVAL = datetime.timedelta(days=3)
 
-app = OpenLLMTyper(help="manage repos")
+app = OpenLLMTyper(help='manage repos')
 
 
-@app.command(name="list", help="list available repo")
+@app.command(name='list', help='list available repo')
 def list_repo(verbose: bool = False):
     if verbose:
         VERBOSE_LEVEL.set(20)
     config = load_config()
     pyaml.pprint(
-        [parse_repo_url(repo, name) for name, repo in config.repos.items()],
-        sort_dicts=False,
-        sort_keys=False,
+        [parse_repo_url(repo, name) for name, repo in config.repos.items()], sort_dicts=False, sort_keys=False
     )
 
 
-@app.command(help="remove given repo")
+@app.command(help='remove given repo')
 def remove(name: str):
     config = load_config()
     if name not in config.repos:
-        output(f"Repo {name} does not exist", style="red")
+        output(f'Repo {name} does not exist', style='red')
         return
 
     del config.repos[name]
     save_config(config)
-    output(f"Repo {name} removed", style="green")
+    output(f'Repo {name} removed', style='green')
 
 
 def _complete_alias(repo_name: str):
     from openllm.model import list_bento
 
     for bento in list_bento(repo_name=repo_name):
-        alias = bento.labels.get("openllm_alias", "").strip()
+        alias = bento.labels.get('openllm_alias', '').strip()
         if alias:
-            for a in alias.split(","):
-                with open(bento.path.parent / a, "w") as f:
+            for a in alias.split(','):
+                with open(bento.path.parent / a, 'w') as f:
                     f.write(bento.version)
 
 
@@ -65,28 +55,15 @@ def _clone_repo(repo: RepoInfo):
     import subprocess
 
     try:
-        subprocess.run(
-            [
-                "git",
-                "clone",
-                "--depth=1",
-                "-b",
-                repo.branch,
-                repo.url,
-                str(repo.path),
-            ],
-            check=True,
-        )
+        subprocess.run(['git', 'clone', '--depth=1', '-b', repo.branch, repo.url, str(repo.path)], check=True)
     except (subprocess.CalledProcessError, FileNotFoundError):
         import dulwich
         import dulwich.porcelain
 
-        dulwich.porcelain.clone(
-            repo.url, str(repo.path), checkout=True, depth=1, branch=repo.branch
-        )
+        dulwich.porcelain.clone(repo.url, str(repo.path), checkout=True, depth=1, branch=repo.branch)
 
 
-@app.command(help="update default repo")
+@app.command(help='update default repo')
 def update():
     config = load_config()
     repos_in_use = set()
@@ -98,59 +75,59 @@ def update():
         repo.path.parent.mkdir(parents=True, exist_ok=True)
         try:
             _clone_repo(repo)
-            output("")
-            output(f"Repo `{repo.name}` updated", style="green")
+            output('')
+            output(f'Repo `{repo.name}` updated', style='green')
         except Exception as e:
             shutil.rmtree(repo.path, ignore_errors=True)
-            output(f"Failed to clone repo {repo.name}", style="red")
+            output(f'Failed to clone repo {repo.name}', style='red')
             output(e)
-    for c in REPO_DIR.glob("*/*/*/*"):
+    for c in REPO_DIR.glob('*/*/*/*'):
         repo_spec = tuple(c.parts[-4:])
         if repo_spec not in repos_in_use:
             shutil.rmtree(c, ignore_errors=True)
-            output(f"Removed unused repo cache {c}")
-    with open(REPO_DIR / "last_update", "w") as f:
+            output(f'Removed unused repo cache {c}')
+    with open(REPO_DIR / 'last_update', 'w') as f:
         f.write(datetime.datetime.now().isoformat())
     for repo_name in config.repos:
         _complete_alias(repo_name)
 
 
 def ensure_repo_updated():
-    last_update_file = REPO_DIR / "last_update"
+    last_update_file = REPO_DIR / 'last_update'
     if not last_update_file.exists():
         if INTERACTIVE.get():
             choice = questionary.confirm(
-                "The repo cache is never updated, do you want to update it to fetch the latest model list?"
+                'The repo cache is never updated, do you want to update it to fetch the latest model list?'
             ).ask()
             if choice:
                 update()
             return
         else:
             output(
-                "The repo cache is never updated, please run `openllm repo update` to fetch the latest model list",
-                style="red",
+                'The repo cache is never updated, please run `openllm repo update` to fetch the latest model list',
+                style='red',
             )
             raise typer.Exit(1)
     last_update = datetime.datetime.fromisoformat(last_update_file.read_text().strip())
     if datetime.datetime.now() - last_update > UPDATE_INTERVAL:
         if INTERACTIVE.get():
             choice = questionary.confirm(
-                "The repo cache is outdated, do you want to update it to fetch the latest model list?"
+                'The repo cache is outdated, do you want to update it to fetch the latest model list?'
             ).ask()
             if choice:
                 update()
         else:
             output(
-                "The repo cache is outdated, please run `openllm repo update` to fetch the latest model list",
-                style="yellow",
+                'The repo cache is outdated, please run `openllm repo update` to fetch the latest model list',
+                style='yellow',
             )
 
 
 GIT_HTTP_RE = re.compile(
-    r"(?P<schema>git|ssh|http|https):\/\/(?P<server>[\.\w\d\-]+)\/(?P<owner>[\w\d\-]+)\/(?P<repo>[\w\d\-\_\.]+)(@(?P<branch>.+))?(\/)?$"
+    r'(?P<schema>git|ssh|http|https):\/\/(?P<server>[\.\w\d\-]+)\/(?P<owner>[\w\d\-]+)\/(?P<repo>[\w\d\-\_\.]+)(@(?P<branch>.+))?(\/)?$'
 )
 GIT_SSH_RE = re.compile(
-    r"git@(?P<server>[\.\w\d-]+):(?P<owner>[\w\d\-]+)\/(?P<repo>[\w\d\-\_\.]+)(@(?P<branch>.+))?(\/)?$"
+    r'git@(?P<server>[\.\w\d-]+):(?P<owner>[\w\d\-]+)\/(?P<repo>[\w\d\-\_\.]+)(@(?P<branch>.+))?(\/)?$'
 )
 
 
@@ -171,27 +148,27 @@ def parse_repo_url(repo_url: str, repo_name: typing.Optional[str] = None) -> Rep
     """
     match = GIT_HTTP_RE.match(repo_url)
     if match:
-        schema = match.group("schema")
+        schema = match.group('schema')
     else:
         match = GIT_SSH_RE.match(repo_url)
         if not match:
-            raise ValueError(f"Invalid git repo url: {repo_url}")
+            raise ValueError(f'Invalid git repo url: {repo_url}')
         schema = None
 
-    if match.group("branch") is not None:
-        repo_url = repo_url[: match.start("branch") - 1]
+    if match.group('branch') is not None:
+        repo_url = repo_url[: match.start('branch') - 1]
 
-    server = match.group("server")
-    owner = match.group("owner")
-    repo = match.group("repo")
-    if repo.endswith(".git"):
+    server = match.group('server')
+    owner = match.group('owner')
+    repo = match.group('repo')
+    if repo.endswith('.git'):
         repo = repo[:-4]
-    branch = match.group("branch") or "main"
+    branch = match.group('branch') or 'main'
 
     if schema is not None:
-        repo_url = f"{schema}://{server}/{owner}/{repo}"
+        repo_url = f'{schema}://{server}/{owner}/{repo}'
     else:
-        repo_url = f"git@{server}:{owner}/{repo}"
+        repo_url = f'git@{server}:{owner}/{repo}'
 
     path = REPO_DIR / server / owner / repo / branch
     return RepoInfo(
@@ -205,40 +182,35 @@ def parse_repo_url(repo_url: str, repo_name: typing.Optional[str] = None) -> Rep
     )
 
 
-@app.command(help="add new repo")
+@app.command(help='add new repo')
 def add(name: str, repo: str):
     name = name.lower()
     if not name.isidentifier():
-        output(
-            f"Invalid repo name: {name}, should only contain letters, numbers and underscores",
-            style="red",
-        )
+        output(f'Invalid repo name: {name}, should only contain letters, numbers and underscores', style='red')
         return
 
     try:
         parse_repo_url(repo)
     except ValueError:
-        output(f"Invalid repo url: {repo}", style="red")
+        output(f'Invalid repo url: {repo}', style='red')
         return
 
     config = load_config()
     if name in config.repos:
-        override = questionary.confirm(
-            f"Repo {name} already exists({config.repos[name]}), override?"
-        ).ask()
+        override = questionary.confirm(f'Repo {name} already exists({config.repos[name]}), override?').ask()
         if not override:
             return
 
     config.repos[name] = repo
     save_config(config)
-    output(f"Repo {name} added", style="green")
+    output(f'Repo {name} added', style='green')
 
 
-@app.command(help="get default repo path")
+@app.command(help='get default repo path')
 def default():
-    output((info := parse_repo_url(load_config().repos["default"], "default")).path)
+    output((info := parse_repo_url(load_config().repos['default'], 'default')).path)
     return info.path
 
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     app()

src/openllm/venv.py

@@ -15,7 +15,7 @@ def _resolve_packages(requirement: typing.Union[pathlib.Path, str]):
     from pip_requirements_parser import RequirementsFile
 
     requirements_txt = RequirementsFile.from_file(str(requirement), include_nested=True)
-    return requirements_txt.requirements
+    return requirements_txt
 
 
 def _filter_preheat_packages(requirements: Iterable) -> list[str]:
@@ -43,14 +43,19 @@ def _resolve_bento_env_specs(bento: BentoInfo):
         lock_file = bento.path / 'env' / 'python' / 'requirements.txt'
 
     reqs = _resolve_packages(lock_file)
-    preheat_packages = _filter_preheat_packages(reqs)
+    packages = reqs.requirements
+    options = reqs.options
+    preheat_packages = _filter_preheat_packages(packages)
     ver = ver_file.read_text().strip()
     return (
         VenvSpec(
             python_version=ver, python_packages=preheat_packages, name_prefix=f"{bento.tag.replace(':', '_')}-1-"
         ),
         VenvSpec(
-            python_version=ver, python_packages=[v.line for v in reqs], name_prefix=f"{bento.tag.replace(':', '_')}-2-"
+            python_version=ver,
+            python_packages=[v.line for v in packages],
+            options=[o.line for o in options],
+            name_prefix=f"{bento.tag.replace(':', '_')}-2-",
         ),
     )
 
@@ -78,7 +83,8 @@ def _ensure_venv(env_spec: VenvSpec, parrent_venv: typing.Optional[pathlib.Path]
                 with open(lib_dir / f'{parrent_venv.name}.pth', 'w+') as f:
                     f.write(str(parent_lib_dir))
             with open(venv / 'requirements.txt', 'w') as f:
-                f.write('\n'.join(sorted(env_spec.python_packages)))
+                with open(venv / 'requirements.txt', 'w') as f:
+                    f.write('\n'.join(env_spec.options + sorted(env_spec.python_packages)))
             run_command(
                 ['python', '-m', 'uv', 'pip', 'install', '-p', str(venv_py), '-r', venv / 'requirements.txt'],
                 silent=VERBOSE_LEVEL.get() < 10,