mirror/OpenLLM
1
0
Fork 0
mirror of https://github.com/bentoml/OpenLLM.git synced 2026-02-20 07:33:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

infra: using ruff formatter (#594)

Browse Source 
Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
This commit is contained in:
Aaron Pham
2023-11-09 12:44:05 -05:00
committed by GitHub
parent 021fd453b9
commit ac377fe490
102 changed files with 5577 additions and 2540 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
openllm-python/src/openllm/testing.py
View File

@@ -1,4 +1,5 @@
"""Tests utilities for OpenLLM."""
from __future__ import annotations
import contextlib
import logging
@@ -9,14 +10,18 @@ import typing as t
import bentoml
import openllm
if t.TYPE_CHECKING:
from openllm_core._typing_compat import LiteralBackend
from openllm_core._typing_compat import LiteralQuantise
logger = logging.getLogger(__name__)
@contextlib.contextmanager
def build_bento(model: str, model_id: str | None = None, quantize: LiteralQuantise | None = None, cleanup: bool = False) -> t.Iterator[bentoml.Bento]:
def build_bento(
model: str, model_id: str | None = None, quantize: LiteralQuantise | None = None, cleanup: bool = False
) -> t.Iterator[bentoml.Bento]:
logger.info('Building BentoML for %s', model)
bento = openllm.build(model, model_id=model_id, quantize=quantize)
yield bento
@@ -24,29 +29,39 @@ def build_bento(model: str, model_id: str | None = None, quantize: LiteralQuanti
logger.info('Deleting %s', bento.tag)
bentoml.bentos.delete(bento.tag)
@contextlib.contextmanager
def build_container(bento: bentoml.Bento | str | bentoml.Tag, image_tag: str | None = None, cleanup: bool = False, **attrs: t.Any) -> t.Iterator[str]:
if isinstance(bento, bentoml.Bento): bento_tag = bento.tag
else: bento_tag = bentoml.Tag.from_taglike(bento)
if image_tag is None: image_tag = str(bento_tag)
def build_container(
bento: bentoml.Bento | str | bentoml.Tag, image_tag: str | None = None, cleanup: bool = False, **attrs: t.Any
) -> t.Iterator[str]:
if isinstance(bento, bentoml.Bento):
bento_tag = bento.tag
else:
bento_tag = bentoml.Tag.from_taglike(bento)
if image_tag is None:
image_tag = str(bento_tag)
executable = shutil.which('docker')
if not executable: raise RuntimeError('docker executable not found')
if not executable:
raise RuntimeError('docker executable not found')
try:
logger.info('Building container for %s', bento_tag)
bentoml.container.build(bento_tag, backend='docker', image_tag=(image_tag,), progress='plain', **attrs,)
bentoml.container.build(bento_tag, backend='docker', image_tag=(image_tag,), progress='plain', **attrs)
yield image_tag
finally:
if cleanup:
logger.info('Deleting container %s', image_tag)
subprocess.check_output([executable, 'rmi', '-f', image_tag])
@contextlib.contextmanager
def prepare(model: str,
model_id: str,
backend: LiteralBackend = 'pt',
deployment_mode: t.Literal['container', 'local'] = 'local',
clean_context: contextlib.ExitStack | None = None,
cleanup: bool = True) -> t.Iterator[str]:
def prepare(
model: str,
model_id: str,
backend: LiteralBackend = 'pt',
deployment_mode: t.Literal['container', 'local'] = 'local',
clean_context: contextlib.ExitStack | None = None,
cleanup: bool = True,
) -> t.Iterator[str]:
if clean_context is None:
clean_context = contextlib.ExitStack()
cleanup = True
@@ -60,4 +75,5 @@ def prepare(model: str,
if deployment_mode == 'container':
container_name = clean_context.enter_context(build_container(bento, image_tag=container_name, cleanup=cleanup))
yield container_name
if cleanup: clean_context.close()
if cleanup:
clean_context.close()