Source code for bentoml._internal.frameworks.picklable_model

from __future__ import annotations

import logging
import typing as t
from types import ModuleType
from typing import TYPE_CHECKING

import cloudpickle

import bentoml
from bentoml import Tag
from bentoml.exceptions import NotFound
from bentoml.models import Model
from bentoml.models import ModelContext

from ..models import PKL_EXT
from ..models import SAVE_NAMESPACE
from ..models.model import ModelSignature
from ..models.model import PartialKwargsModelOptions as ModelOptions

if TYPE_CHECKING:
    ModelType = t.Any

MODULE_NAME = "bentoml.picklable_model"
API_VERSION = "v1"


logger = logging.getLogger(__name__)


[docs] def get(tag_like: str | Tag) -> Model: model = bentoml.models.get(tag_like) if model.info.module not in (MODULE_NAME, __name__): raise NotFound( f"Model {model.tag} was saved with module {model.info.module}, not loading with {MODULE_NAME}." ) return model
[docs] def load_model(bento_model: str | Tag | Model) -> ModelType: """ Load the picklable model with the given tag from the local BentoML model store. Args: bento_model: Either the tag of the model to get from the store, or a BentoML :class:`~bentoml.Model` instance to load the model from. Returns: The picklable model loaded from the model store or BentoML :obj:`~bentoml.Model`. Example: .. code-block:: python import bentoml picklable_model = bentoml.picklable_model.load_model('my_model:latest') """ # noqa if not isinstance(bento_model, Model): bento_model = get(bento_model) if bento_model.info.module not in (MODULE_NAME, __name__): raise NotFound( f"Model {bento_model.tag} was saved with module {bento_model.info.module}, not loading with {MODULE_NAME}." ) model_file = bento_model.path_of(f"{SAVE_NAMESPACE}{PKL_EXT}") with open(model_file, "rb") as f: # The protocol version used is detected automatically, so we do not # have to specify it. return cloudpickle.load(f)
[docs] def save_model( name: Tag | str, model: ModelType, *, signatures: dict[str, ModelSignature] | None = None, labels: t.Dict[str, str] | None = None, custom_objects: t.Dict[str, t.Any] | None = None, external_modules: t.List[ModuleType] | None = None, metadata: t.Dict[str, t.Any] | None = None, ) -> bentoml.Model: """ Save a model instance to BentoML modelstore. Args: name: Name for given model instance. This should pass Python identifier check. model: Instance of model to be saved. signatures: Methods to expose for running inference on the target model. Signatures are used for creating Runner instances when serving model with bentoml.Service labels: user-defined labels for managing models, e.g. team=nlp, stage=dev custom_objects: user-defined additional python objects to be saved alongside the model, e.g. a tokenizer instance, preprocessor function, model configuration json external_modules: user-defined additional python modules to be saved alongside the model or custom objects, e.g. a tokenizer module, preprocessor module, model configuration module metadata: Custom metadata for given model. Returns: :obj:`~bentoml.Tag`: A :obj:`tag` with a format ``name:version`` where ``name`` is the user-defined model's name, and a generated ``version``. Examples: .. code-block:: python import bentoml bento_model = bentoml.picklable.save_model('picklable_pyobj', model) """ # noqa context = ModelContext( framework_name="cloudpickle", framework_versions={"cloudpickle": cloudpickle.__version__}, ) if signatures is None: signatures = {"__call__": ModelSignature(batchable=False)} logger.info( 'Using the default model signature for pickable model (%s) for model "%s".', signatures, name, ) with bentoml.models._create( # type: ignore name, module=MODULE_NAME, api_version=API_VERSION, labels=labels, custom_objects=custom_objects, external_modules=external_modules, metadata=metadata, context=context, signatures=signatures, options=ModelOptions(), ) as bento_model: with open(bento_model.path_of(f"{SAVE_NAMESPACE}{PKL_EXT}"), "wb") as f: cloudpickle.dump(model, f) return bento_model
def get_runnable(bento_model: Model): """ Private API: use :obj:`~bentoml.Model.to_runnable` instead. """ partial_kwargs: t.Dict[str, t.Any] = bento_model.info.options.partial_kwargs # type: ignore class PicklableRunnable(bentoml.Runnable): SUPPORTED_RESOURCES = ("cpu",) SUPPORTS_CPU_MULTI_THREADING = False def __init__(self): super().__init__() self.model = load_model(bento_model) def _get_run(method_name: str, partial_kwargs: t.Dict[str, t.Any] | None = None): if partial_kwargs is None: partial_kwargs = {} def _run(self: PicklableRunnable, *args: t.Any, **kwargs: t.Any) -> t.Any: assert isinstance(method_name, str), repr(method_name) return getattr(self.model, method_name)( *args, **dict(partial_kwargs, **kwargs) ) return _run for method_name, options in bento_model.info.signatures.items(): assert isinstance(method_name, str), repr(method_name) method_partial_kwargs = partial_kwargs.get(method_name) PicklableRunnable.add_method( _get_run(method_name, method_partial_kwargs), name=method_name, batchable=options.batchable, batch_dim=options.batch_dim, input_spec=options.input_spec, output_spec=options.output_spec, ) return PicklableRunnable