caikit.core.model_manager

Most logic interacting with models. Can load, etc.

Classes

ModelManager

Manage the models or resources for library.

Functions

get_valid_module_ids()

Get a dictionary mapping all module IDs to the string names of the

Module Contents

caikit.core.model_manager.get_valid_module_ids()

Get a dictionary mapping all module IDs to the string names of the implementing classes.

class caikit.core.model_manager.ModelManager

Manage the models or resources for library.

_singleton_module_cache

_trainers

_finders

_job_predictors

_initializers

__singleton_lock

initialize_components()

Proactively initialize all configured trainer/finder/initializer component instances. This is a separate call to enable explicit config.

train(module: Type[caikit.core.modules.base.ModuleBase] | str, *args, trainer: str | caikit.core.model_management.ModelTrainerBase = 'default', save_path: str | caikit.interfaces.common.data_model.stream_sources.S3Path | None = None, save_with_id: bool = False, model_name: str | None = None, wait: bool = False, **kwargs) → caikit.core.model_management.ModelTrainerFutureBase

Train an instance of the given module with the given args and kwargs using the given trainer.

Each module’s train function encapsulates the code needed to perform the training locally. This top-level train function provides the wrapper functionality to delegate the execution of the module’s train function to an alternate framework using a ModelTrainerBase. It also allows training to be launched asynchronously.

Args:

module (Union[Type[ModuleBase], str]): The module class or guid for

the module to train

*args: Additional positional args to pass through to the module’s

train function

Kwargs:

trainer (Union[str, ModelTrainerBase]): The trainer to use. If given

as a string, this is a key in the global config at model_management.trainers.

save_path (Optional[Union[str, S3Path]]): Base path where the model should be

saved (may be relative to a remote trainer’s filesystem, or link to S3 storage)

save_with_id (bool): Inject the training ID into the save path for

the output model

model_name (Optional[str]): Name of model that will be appended

to the end of the save_path

wait (bool): Wait for training to complete before returning **kwargs: Additional keyword arguments to pass through to the

modules’s train function

Returns:

model_future (ModelFutureBase): The future handle

to the model which holds the status of the in-flight training.

start_prediction_job(model: caikit.core.modules.base.ModuleBase, prediction_func_name: str, *args, predictor: str | caikit.core.model_management.JobPredictorBase = 'default', wait: bool = False, **kwargs) → caikit.core.model_management.JobPredictorFutureBase

Start a prediction job using a job_predictor.

Args:

model (ModuleBase): Loaded model to run prediction on prediction_func_name (str): String reference to name of function to run predictor (Union[str, JobPredictorBase], optional): Which job_predictor to use.

Defaults to “default”.

wait (bool, optional): Weather to wait for job to finish. Defaults to False.

Returns:

JobPredictorFutureBase: Future to track job result

get_model_future(training_id: str) → caikit.core.model_management.ModelTrainerFutureBase

Get the future handle to an in-progress training

Args:

training_id (str): The ID string from the original training

submission’s ModelFuture

Returns:

model_future (ModelTrainerFutureBase): The future handle

to the model which holds the status of the in-flight training.

get_prediction_future(prediction_id: str) → caikit.core.model_management.JobPredictorFutureBase

Get the future handle to an in-progress prediction job

Args:

prediction_id (str): The ID string from the original prediction

submission’s ModelFuture

Returns:

prediction_future (JobPredictorFutureBase): The future handle

to the job which holds the status of the in-flight prediction.

load(module_path: str | io.BytesIO | bytes, *, load_singleton: bool = False, finder: str | caikit.core.model_management.ModelFinderBase = 'default', initializer: str | caikit.core.model_management.ModelInitializerBase = 'default', **kwargs)

Load a model and return an instantiated object on which we can run inference.

Args:

module_path (str | BytesIO | bytes): A module path (identifier) to

one of the following: 1. A directory containing a yaml config file in the top level. 2. A zip archive containing either a yaml config file in the

top level when extracted, or a directory containing a yaml config file in the top level.

3. A BytesIO object corresponding to a zip archive containing

   either a yaml config file in the top level when extracted, or a directory containing a yaml config file in the top level.

4. A bytes object corresponding to a zip archive containing

   either a yaml config file in the top level when extracted, or a directory containing a yaml config file in the top level.

5. A string that is understood by the configured

   finder/initializer

Kwargs:

load_singleton (bool): Load this model as a singleton finder (Union[str, ModelFinderBase]): Finder to use when loading

this model. If passed as a string, this names the finder in the global config model_management.finders section.

initializer (Union[str, ModelInitializerBase]): Loader to use when

initializint this model. If passed as a string, this is the name of the initializer in the global config model_management.initializers section.

Returns:

model (ModuleBase) Model object that is loaded, configured, and

ready for prediction.

extract(zip_path: str, model_path: str, force_overwrite: bool = False) → str

Method to extract a downloaded archive to a specified directory.

Args:

zip_path (str): Location of .zip file to extract. model_path (str): Model directory where the archive should be

unzipped unzipped.

force_overwrite: bool (Defaults to false)

Force an overwrite to model_path, even if the folder exists

Returns:

str: Output path where the model archive is unzipped.

resolve_and_load(path_or_name_or_model_reference: str | caikit.core.modules.base.ModuleBase, **kwargs)

Try our best to load a model, given a path or a name. Simply returns any loaded model passed in. This exists to ease the burden on workflow developers who need to accept individual modules in their API, where users may have references to custom models or may only have the ability to give the name of a stock model.

Args:

path_or_name_or_model_reference (str, ModuleBase): Either a

• Path to a model on disk

• Name of a model that the catalog knows about

• Loaded module

**kwargs: Any keyword arguments to pass along to ModelManager.load()

or ModelManager.download()

e.g. parent_dir

Returns:

A loaded module

Examples:

>>> stock_syntax_model = manager.resolve_and_load('syntax_izumo_en_stock')
>>> local_categories_model = manager.resolve_and_load('path/to/categories/model')
>>> some_custom_model = manager.resolve_and_load(some_custom_model)

get_singleton_model_cache_info()

Returns information about the singleton cache in {hash: module type} format

Returns:

Dict[str, type]: A dictionary of model hashes to model types

clear_singleton_cache()

Clears the cache of singleton models. Useful to release references of models, as long as you know that they are no longer held elsewhere and you won’t be loading them again.

Returns:

None

get_trainer(trainer: str | caikit.core.model_management.ModelTrainerBase) → caikit.core.model_management.ModelTrainerBase

Get the configured model trainer or the one passed by value

get_finder(finder: str | caikit.core.model_management.ModelFinderBase) → caikit.core.model_management.ModelFinderBase

Get the configured model finder or the one passed by value

get_initializer(initializer: str | caikit.core.model_management.ModelInitializerBase) → caikit.core.model_management.ModelInitializerBase

Get the configured model initializer or the one passed by value

get_predictor(inferencer: str | caikit.core.model_management.JobPredictorBase) → caikit.core.model_management.JobPredictorBase

Get the configured job predictor or the one passed by value

get_module_backends(initialize: bool = True) → List[caikit.core.module_backends.base.BackendBase]

Convenience method to get access to the configured module backends if any have been configured

Args:

initialize (bool): Initialize the components from config

Returns:

backends (List[BackendBase]): The list of backend instances that

have been configured

_do_load(module_path, load_singleton, finder, initializer, **kwargs)

Load a model from a directory.

Args:

module_path (str): Path to directory. At the top level of directory

is config.yml which holds info about the model.

load_singleton (bool): Load this model as a singleton finder (Union[str, ModelFinderBase]): Finder to use when loading

this model. If passed as a string, this names the finder in the global config model_management.finders section.

initializer (Union[str, ModelInitializerBase]): Loader to use when

loading this model. If passed as a string, this is the name of the initializer in the global config model_management.initializers section.

Returns:

subclass of caikit.core.modules.ModuleBase: Model object that is

loaded, configured, and ready for prediction.

_load_from_zipfile(module_path, load_singleton, finder, initializer, **kwargs)

Load a model from a zip archive.

Args:

module_path (str): Path to directory. At the top level of directory

is config.yml which holds info about the model.

load_singleton (bool): Load this model as a singleton finder (Union[str, ModelFinderBase]): Finder to use when loading

this model. If passed as a string, this names the finder in the global config model_management.finders section.

initializer (Union[str, ModelInitializerBase]): Loader to use when

loading this model. If passed as a string, this is the name of the initializer in the global config model_management.initializers section.

Returns:

subclass of caikit.core.modules.ModuleBase: Model object that is

loaded, configured, and ready for prediction.

_singleton_lock(load_singleton: bool)

Helper contextmanager that will only lock the singleton cache if this load is a singleton load

static _get_component(component: str | caikit.core.toolkit.factory.FactoryConstructible, component_dict: Dict[str, caikit.core.toolkit.factory.FactoryConstructible], component_factory: caikit.core.toolkit.factory.Factory, component_name: str, component_cfg: dict, component_type: type) → caikit.core.toolkit.factory.FactoryConstructible

Common logic for resolving components from config

NOTE: This is done lazily to avoid relying on import order and to allow

for dynamic config changes