mirror of
https://github.com/bentoml/OpenLLM.git
synced 2026-02-19 23:24:12 -05:00
4a6f13ddd2
* feat(type): provide client stubs separation of concern for more brevity code base Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com> * docs: update changelog Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com> --------- Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
from __future__ import annotations
|
|
import contextlib
|
|
import logging
|
|
import shutil
|
|
import subprocess
|
|
import typing as t
|
|
|
|
import bentoml
|
|
import openllm
|
|
|
|
if t.TYPE_CHECKING:
|
|
from openllm_core._typing_compat import LiteralBackend, 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]:
|
|
logger.info('Building BentoML for %s', model)
|
|
bento = openllm.build(model, model_id=model_id, quantize=quantize)
|
|
yield bento
|
|
if cleanup:
|
|
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)
|
|
executable = shutil.which('docker')
|
|
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)
|
|
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]:
|
|
if clean_context is None:
|
|
clean_context = contextlib.ExitStack()
|
|
cleanup = True
|
|
llm = openllm.LLM[t.Any, t.Any](model_id, backend=backend)
|
|
bento_tag = bentoml.Tag.from_taglike(f'{llm.llm_type}-service:{llm.tag.version}')
|
|
if not bentoml.list(bento_tag):
|
|
bento = clean_context.enter_context(build_bento(model, model_id=model_id, cleanup=cleanup))
|
|
else:
|
|
bento = bentoml.get(bento_tag)
|
|
container_name = f'openllm-{model}-{llm.llm_type}'.replace('-', '_')
|
|
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()
|