Source code for bentoml.models

from __future__ import annotations

import typing as t
from contextlib import contextmanager
from types import ModuleType
from typing import TYPE_CHECKING

from simple_di import Provide
from simple_di import inject

from ._internal.configuration.containers import BentoMLContainer
from ._internal.models import Model
from ._internal.models import ModelContext
from ._internal.models import ModelOptions
from ._internal.tag import Tag
from ._internal.utils import calc_dir_size
from ._internal.utils.analytics import ModelSaveEvent
from ._internal.utils.analytics import track
from .exceptions import BentoMLException

if TYPE_CHECKING:
    from ._internal.cloud import BentoCloudClient
    from ._internal.models import ModelStore
    from ._internal.models.model import ModelSignaturesType


[docs]@inject def list( # pylint: disable=redefined-builtin tag: t.Optional[t.Union[Tag, str]] = None, *, _model_store: "ModelStore" = Provide[BentoMLContainer.model_store], ) -> t.List["Model"]: return _model_store.list(tag)
[docs]@inject def get( tag: t.Union[Tag, str], *, _model_store: "ModelStore" = Provide[BentoMLContainer.model_store], model_aliases: t.Dict[str, str] = Provide[BentoMLContainer.model_aliases], ) -> "Model": if isinstance(tag, str) and tag in model_aliases: tag = model_aliases[tag] return _model_store.get(tag)
[docs]@inject def delete( tag: t.Union[Tag, str], *, _model_store: "ModelStore" = Provide[BentoMLContainer.model_store], ): _model_store.delete(tag)
[docs]@inject def import_model( path: str, input_format: t.Optional[str] = None, *, protocol: t.Optional[str] = None, user: t.Optional[str] = None, passwd: t.Optional[str] = None, params: t.Optional[t.Dict[str, str]] = None, subpath: t.Optional[str] = None, _model_store: "ModelStore" = Provide[BentoMLContainer.model_store], ) -> Model: """ Import a bento model exported with :code:`bentoml.models.export_model`. To import a model saved with a framework, see the :code:`save` function under the relevant framework, e.g. :code:`bentoml.sklearn.save`. Examples: .. code-block:: python # imports 'my_model' from '/path/to/folder/my_model.bentomodel' bentoml.models.import_model('/path/to/folder/my_model.bentomodel') # imports 'my_model' from '/path/to/folder/my_model.tar.gz' # currently supported formats are tar.gz ('gz'), tar.xz ('xz'), tar.bz2 ('bz2'), and zip bentoml.models.import_model('/path/to/folder/my_model.tar.gz') # treats 'my_model.ext' as a gzipped tarfile bentoml.models.import_model('/path/to/folder/my_model.ext', 'gz') # imports 'my_model', which is stored as an uncompressed folder, from '/path/to/folder/my_model/' bentoml.models.import_model('/path/to/folder/my_model', 'folder') # imports 'my_model' from the S3 bucket 'my_bucket', path 'folder/my_model.bentomodel' # requires `fs-s3fs <https://pypi.org/project/fs-s3fs/>`_ ('pip install fs-s3fs') bentoml.models.import_model('s3://my_bucket/folder/my_model.bentomodel') bentoml.models.import_model('my_bucket/folder/my_model.bentomodel', protocol='s3') bentoml.models.import_model('my_bucket', protocol='s3', subpath='folder/my_model.bentomodel') bentoml.models.import_model('my_bucket', protocol='s3', subpath='folder/my_model.bentomodel', user='<AWS access key>', passwd='<AWS secret key>', params={'acl': 'public-read', 'cache-control': 'max-age=2592000,public'}) For a more comprehensive description of what each of the keyword arguments (:code:`protocol`, :code:`user`, :code:`passwd`, :code:`params`, and :code:`subpath`) mean, see the `FS URL documentation <https://docs.pyfilesystem.org/en/latest/openers.html>`_. Args: tag: the tag of the model to export path: can be one of two things: * a folder on the local filesystem * an `FS URL <https://docs.pyfilesystem.org/en/latest/openers.html>`_, for example :code:`'s3://my_bucket/folder/my_model.bentomodel'` protocol: (expert) The FS protocol to use when exporting. Some example protocols are :code:`'ftp'`, :code:`'s3'`, and :code:`'userdata'` user: (expert) the username used for authentication if required, e.g. for FTP passwd: (expert) the username used for authentication if required, e.g. for FTP params: (expert) a map of parameters to be passed to the FS used for export, e.g. :code:`{'proxy': 'myproxy.net'}` for setting a proxy for FTP subpath: (expert) the path inside the FS that the model should be exported to _model_store: the model store to save the model to Returns: Model: the imported model """ return Model.import_from( path, input_format, protocol=protocol, user=user, passwd=passwd, params=params, subpath=subpath, ).save(_model_store)
[docs]@inject def export_model( tag: t.Union[Tag, str], path: str, output_format: t.Optional[str] = None, *, protocol: t.Optional[str] = None, user: t.Optional[str] = None, passwd: t.Optional[str] = None, params: t.Optional[t.Dict[str, str]] = None, subpath: t.Optional[str] = None, _model_store: "ModelStore" = Provide[BentoMLContainer.model_store], ) -> str: """ Export a BentoML model. Examples: .. code-block:: python # exports 'my_model' to '/path/to/folder/my_model-version.bentomodel' in BentoML's default format bentoml.models.export_model('my_model:latest', '/path/to/folder') # note that folders can only be passed if exporting to the local filesystem; otherwise the # full path, including the desired filename, must be passed # exports 'my_model' to '/path/to/folder/my_model.bentomodel' in BentoML's default format bentoml.models.export_model('my_model:latest', '/path/to/folder/my_model') bentoml.models.export_model('my_model:latest', '/path/to/folder/my_model.bentomodel') # exports 'my_model' to '/path/to/folder/my_model.tar.gz in gzip format # currently supported formats are tar.gz ('gz'), tar.xz ('xz'), tar.bz2 ('bz2'), and zip bentoml.models.export_model('my_model:latest', '/path/to/folder/my_model.tar.gz') bentoml.models.export_model('my_model:latest', '/path/to/folder/my_model.tar.gz', 'gz') # exports 'my_model' to '/path/to/folder/my_model/ as a folder bentoml.models.export_model('my_model:latest', '/path/to/folder/my_model', 'folder') # exports 'my_model' to the S3 bucket 'my_bucket' as 'folder/my_model-version.bentomodel' bentoml.models.export_model('my_model:latest', 's3://my_bucket/folder') bentoml.models.export_model('my_model:latest', 'my_bucket/folder', protocol='s3') bentoml.models.export_model('my_model:latest', 'my_bucket', protocol='s3', subpath='folder') bentoml.models.export_model('my_model:latest', 'my_bucket', protocol='s3', subpath='folder', user='<AWS access key>', passwd='<AWS secret key>', params={'acl': 'public-read', 'cache-control': 'max-age=2592000,public'}) For a more comprehensive description of what each of the keyword arguments (:code:`protocol`, :code:`user`, :code:`passwd`, :code:`params`, and :code:`subpath`) mean, see the `FS URL documentation <https://docs.pyfilesystem.org/en/latest/openers.html>`_. Args: tag: the tag of the model to export path: can be one of two things: * a folder on the local filesystem * an `FS URL <https://docs.pyfilesystem.org/en/latest/openers.html>`_, for example, :code:`'s3://my_bucket/folder/my_model.bentomodel'` protocol: (expert) The FS protocol to use when exporting. Some example protocols are :code:`'ftp'`, :code:`'s3'`, and :code:`'userdata'`. user: (expert) the username used for authentication if required, e.g. for FTP passwd: (expert) the username used for authentication if required, e.g. for FTP params: (expert) a map of parameters to be passed to the FS used for export, e.g. :code:`{'proxy': 'myproxy.net'}` for setting a proxy for FTP subpath: (expert) the path inside the FS that the model should be exported to _model_store: the model store to get the model to save from Returns: str: A representation of the path that the model was exported to. If it was exported to the local filesystem, this will be the OS path to the exported model. Otherwise, it will be an FS URL. """ model = get(tag, _model_store=_model_store) return model.export( path, output_format, protocol=protocol, user=user, passwd=passwd, params=params, subpath=subpath, )
[docs]@inject def push( tag: t.Union[Tag, str], *, force: bool = False, _model_store: "ModelStore" = Provide[BentoMLContainer.model_store], _cloud_client: BentoCloudClient = Provide[BentoMLContainer.bentocloud_client], ): model_obj = _model_store.get(tag) if not model_obj: raise BentoMLException(f"Model {tag} not found in local store") _cloud_client.push_model(model_obj, force=force)
[docs]@inject def pull( tag: t.Union[Tag, str], *, force: bool = False, _cloud_client: BentoCloudClient = Provide[BentoMLContainer.bentocloud_client], ) -> Model: return _cloud_client.pull_model(tag, force=force)
if t.TYPE_CHECKING: class CreateKwargs(t.TypedDict): module: str api_version: str | None signatures: t.Required[ModelSignaturesType] labels: dict[str, t.Any] | None options: ModelOptions | None custom_objects: dict[str, t.Any] | None external_modules: t.List[ModuleType] | None metadata: dict[str, t.Any] | None context: t.Required[ModelContext] _model_store: t.NotRequired[ModelStore] @t.overload @contextmanager def _create( name: Tag | str, /, **attrs: t.Unpack[CreateKwargs] ) -> t.Generator[Model, None, None]: ... @t.overload @contextmanager def _create( name: Tag | str, *, module: str = ..., api_version: str | None = ..., signatures: ModelSignaturesType, labels: dict[str, t.Any] | None = ..., options: ModelOptions | None = ..., custom_objects: dict[str, t.Any] | None = ..., external_modules: t.List[ModuleType] | None = ..., metadata: dict[str, t.Any] | None = ..., context: ModelContext, _model_store: ModelStore = ..., ) -> t.Generator[Model, None, None]: ... @inject @contextmanager def _create( name: Tag | str, *, module: str = "", api_version: str | None = None, signatures: ModelSignaturesType | None = None, labels: dict[str, t.Any] | None = None, options: ModelOptions | None = None, custom_objects: dict[str, t.Any] | None = None, external_modules: t.List[ModuleType] | None = None, metadata: dict[str, t.Any] | None = None, context: ModelContext | None = None, _model_store: ModelStore = Provide[BentoMLContainer.model_store], ) -> t.Generator[Model, None, None]: options = ModelOptions() if options is None else options api_version = "v1" if api_version is None else api_version res = Model.create( name, module=module, api_version=api_version, labels=labels, signatures=signatures or {}, options=options, custom_objects=custom_objects, metadata=metadata, context=context or ModelContext("", {}), ) external_modules = [] if external_modules is None else external_modules imported_modules: t.List[ModuleType] = [] try: res.enter_cloudpickle_context(external_modules, imported_modules) yield res except Exception: raise else: res.flush() res.save(_model_store) track( ModelSaveEvent( module=res.info.module, model_size_in_kb=calc_dir_size(res.path_of("/")) / 1024, ), ) finally: res.exit_cloudpickle_context(imported_modules)
[docs]@inject @contextmanager def create( name: Tag | str, *, labels: dict[str, t.Any] | None = None, metadata: dict[str, t.Any] | None = None, _model_store: ModelStore = Provide[BentoMLContainer.model_store], module: str = "", # deprecated api_version: str = "v1", # deprecated signatures: ModelSignaturesType = {}, # deprecated options: ModelOptions = ModelOptions(), # deprecated custom_objects: dict[str, t.Any] | None = None, # deprecated external_modules: t.List[ModuleType] | None = None, # deprecated context: ModelContext = ModelContext( framework_name="", framework_versions={} ), # deprecated ) -> t.Generator[Model, None, None]: """ Create a model instance in BentoML modelstore. The following function arguments are deprecated: module, api_version, signatures, options, custom_objects, external_modules, context """ res = Model.create( name, module=module, api_version=api_version, labels=labels, signatures=signatures, options=options, custom_objects=custom_objects, metadata=metadata, context=context, ) external_modules = [] if external_modules is None else external_modules imported_modules: t.List[ModuleType] = [] try: res.enter_cloudpickle_context(external_modules, imported_modules) yield res except Exception: raise else: res.flush() res.save(_model_store) track( ModelSaveEvent( module=res.info.module, model_size_in_kb=calc_dir_size(res.path_of("/")) / 1024, ), ) finally: res.exit_cloudpickle_context(imported_modules)
__all__ = [ "list", "get", "delete", "import_model", "export_model", "push", "pull", "ModelContext", "ModelOptions", "create", ]