diff --git a/caikit/config/config.yml b/caikit/config/config.yml index c3ca32c55..249dd0ad2 100644 --- a/caikit/config/config.yml +++ b/caikit/config/config.yml @@ -51,6 +51,14 @@ model_management: # List of module backend configurations in priority order backend_priority: - type: LOCAL + loaders: + default: + type: CORE + config: {} + sizers: + default: + type: MODEL_MESH + config: {} log: # Default level for all python loggers diff --git a/caikit/runtime/model_management/core_model_loader.py b/caikit/runtime/model_management/core_model_loader.py new file mode 100644 index 000000000..deaf1b2c6 --- /dev/null +++ b/caikit/runtime/model_management/core_model_loader.py @@ -0,0 +1,62 @@ +# Copyright The Caikit Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Standard +from typing import Optional, Union + +# Third Party +from prometheus_client import Summary + +# First Party +import alog + +# Local +from caikit.core import MODEL_MANAGER, ModuleBase +from caikit.core.model_management import ModelFinderBase, ModelInitializerBase +from caikit.runtime.model_management.model_loader_base import ModelLoaderBase + +log = alog.use_channel("MODEL-LOADER") + +CAIKIT_CORE_LOAD_DURATION_SUMMARY = Summary( + "caikit_core_load_model_duration_seconds", + "Summary of the duration (in seconds) of caikit.core.load(model)", + ["model_type"], +) + + +class CoreModelLoader(ModelLoaderBase): + """The CoreModelLoader loads a model using the caikit core.ModelManager""" + + name = "CORE" + + def load_module_instance( + self, + model_path: str, + model_id: str, + model_type: str, + finder: Optional[Union[str, ModelFinderBase]] = None, + initializer: Optional[Union[str, ModelInitializerBase]] = None, + ) -> ModuleBase: + """Start loading a model from disk and associate the ID/size with it""" + log.info("", "Loading model '%s'", model_id) + + # Only pass finder/initializer if they have values so that defaults are used otherwise + load_kwargs = {} + if finder: + load_kwargs["finder"] = finder + if initializer: + load_kwargs["initializer"] = initializer + + # Load using the caikit.core + with CAIKIT_CORE_LOAD_DURATION_SUMMARY.labels(model_type=model_type).time(): + return MODEL_MANAGER.load(model_path, **load_kwargs) diff --git a/caikit/runtime/model_management/directory_model_sizer.py b/caikit/runtime/model_management/directory_model_sizer.py new file mode 100644 index 000000000..1b5d8c733 --- /dev/null +++ b/caikit/runtime/model_management/directory_model_sizer.py @@ -0,0 +1,88 @@ +# Copyright The Caikit Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Standard +from pathlib import Path +from typing import Dict +import os + +# Third Party +import grpc + +# First Party +import aconfig +import alog + +# Local +from caikit.runtime.model_management.model_sizer_base import ModelSizerBase +from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException + +log = alog.use_channel("DIRECTORY-SIZER") + + +class DirectoryModelSizer(ModelSizerBase): + """DirectoryModelSizer. This class calculates a models size based on the + size of the files in the model directory + + ! Note: It caches the size of the directory after first sizing which can cause + race conditions in certain situations. + """ + + name = "DIRECTORY" + + def __init__(self, config: aconfig.Config, instance_name: str): + super().__init__(config, instance_name) + # Cache of archive sizes: directory model path -> archive size in bytes + self.model_directory_size: Dict[str, int] = {} + + def get_model_size(self, model_id, local_model_path, model_type) -> int: + """ + Returns the estimated memory footprint of a model + Args: + model_id: The model identifier, used for informative logging + cos_model_path: The path to the model archive in S3 storage + model_type: The type of model, used to adjust the memory estimate + Returns: + The estimated size in bytes of memory that would be used by loading this model + """ + # Return the cached model size if one exists + if model_size := self.model_directory_size.get(local_model_path): + return model_size + + # Calculate the model size and add it to the cache. This uses last in + # methodology so that the most recent size is used during parallel access + dir_size = self.__get_directory_size(model_id, local_model_path) + self.model_directory_size[local_model_path] = dir_size + return dir_size + + def __get_directory_size(self, model_id, local_model_path) -> int: + """Get the size of a directory""" + try: + if os.path.isdir(local_model_path): + # Walk the directory to size all files + return sum( + file.stat().st_size + for file in Path(local_model_path).rglob("*") + if file.is_file() + ) + + # Probably just an archive file + return os.path.getsize(local_model_path) + except FileNotFoundError as ex: + message = ( + f"Failed to estimate size of model '{model_id}'," + f"file '{local_model_path}' not found" + ) + log.error("", message) + raise CaikitRuntimeException(grpc.StatusCode.NOT_FOUND, message) from ex diff --git a/caikit/runtime/model_management/factories.py b/caikit/runtime/model_management/factories.py new file mode 100644 index 000000000..ff637e5cb --- /dev/null +++ b/caikit/runtime/model_management/factories.py @@ -0,0 +1,33 @@ +# Copyright The Caikit Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Global factories for model management +""" + +# Local +from caikit.core.toolkit.factory import ImportableFactory +from caikit.runtime.model_management.core_model_loader import CoreModelLoader +from caikit.runtime.model_management.directory_model_sizer import DirectoryModelSizer +from caikit.runtime.model_management.mm_model_sizer import ModelMeshModelSizer + +# Model Loader factory. A loader is responsible for constructing +# a LoadedModel instance +model_loader_factory = ImportableFactory("ModelLoader") +model_loader_factory.register(CoreModelLoader) + +# Model Sizer factory. A sizer is responsible for estimating +# the size of a model +model_sizer_factory = ImportableFactory("ModelSizer") +model_sizer_factory.register(DirectoryModelSizer) +model_sizer_factory.register(ModelMeshModelSizer) diff --git a/caikit/runtime/model_management/mm_model_sizer.py b/caikit/runtime/model_management/mm_model_sizer.py new file mode 100644 index 000000000..5d305efc2 --- /dev/null +++ b/caikit/runtime/model_management/mm_model_sizer.py @@ -0,0 +1,71 @@ +# Copyright The Caikit Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# First Party +import alog + +# Local +from caikit import get_config +from caikit.runtime.model_management.directory_model_sizer import DirectoryModelSizer + +log = alog.use_channel("MM-SIZER") + + +class ModelMeshModelSizer(DirectoryModelSizer): + """ModelMeshModelSizer. This class estimates a models size based on + the contents of the directory multiplied by a model specific + constant""" + + name = "MODEL_MESH" + + def get_model_size(self, model_id, local_model_path, model_type) -> int: + """ + Returns the estimated memory footprint of a model + Args: + model_id: The model identifier, used for informative logging + cos_model_path: The path to the model archive in S3 storage + model_type: The type of model, used to adjust the memory estimate + Returns: + The estimated size in bytes of memory that would be used by loading this model + """ + + if ( + model_type + in get_config().inference_plugin.model_mesh.model_size_multipliers + ): + multiplier = ( + get_config().inference_plugin.model_mesh.model_size_multipliers[ + model_type + ] + ) + log.debug( + "Using size multiplier '%f' for model '%s' to estimate model size", + multiplier, + model_id, + ) + else: + multiplier = ( + get_config().inference_plugin.model_mesh.default_model_size_multiplier + ) + log.info( + "", + "No configured model size multiplier found for model type '%s' for model '%s'. " + "Using default multiplier '%f'", + model_type, + model_id, + multiplier, + ) + return int( + super().get_model_size(model_id, local_model_path, model_type) * multiplier + ) diff --git a/caikit/runtime/model_management/model_loader.py b/caikit/runtime/model_management/model_loader_base.py similarity index 74% rename from caikit/runtime/model_management/model_loader.py rename to caikit/runtime/model_management/model_loader_base.py index 1a1ce043d..4421c0517 100644 --- a/caikit/runtime/model_management/model_loader.py +++ b/caikit/runtime/model_management/model_loader_base.py @@ -15,46 +15,70 @@ from concurrent.futures import ThreadPoolExecutor from functools import partial from typing import Callable, Optional, Union +import abc # Third Party from grpc import StatusCode -from prometheus_client import Summary # First Party +import aconfig import alog # Local from caikit.config import get_config from caikit.core import MODEL_MANAGER, ModuleBase from caikit.core.model_management import ModelFinderBase, ModelInitializerBase +from caikit.core.toolkit.factory import FactoryConstructible from caikit.runtime.model_management.batcher import Batcher from caikit.runtime.model_management.loaded_model import LoadedModel from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException log = alog.use_channel("MODEL-LOADER") -CAIKIT_CORE_LOAD_DURATION_SUMMARY = Summary( - "caikit_core_load_model_duration_seconds", - "Summary of the duration (in seconds) of caikit.core.load(model)", - ["model_type"], -) +class ModelLoaderBase(FactoryConstructible): + """Model Loader Base class which describes how models are loaded.""" -class ModelLoader: - """Model Loader class. The singleton class contains the core implementation details - for loading models in from S3.""" + _load_thread_pool = None - __instance = None + def __init__(self, config: aconfig.Config, instance_name: str): + """A FactoryConstructible object must be constructed with a config + object that it uses to pull in all configuration + """ + if ModelLoaderBase._load_thread_pool is None: + ModelLoaderBase._load_thread_pool = ThreadPoolExecutor( + get_config().runtime.load_threads + ) - def __init__(self): - # Re-instantiating this is a programming error - assert self.__class__.__instance is None, "This class is a singleton!" - ModelLoader.__instance = self - self._load_thread_pool = ThreadPoolExecutor(get_config().runtime.load_threads) + super().__init__(config, instance_name) # Instead of storing config-based batching information here, we call # get_config() when needed to support dynamic config changes for # batching + @abc.abstractmethod + def load_module_instance( + self, + model_path: str, + model_id: str, + model_type: str, + finder: Optional[Union[str, ModelFinderBase]] = None, + initializer: Optional[Union[str, ModelInitializerBase]] = None, + ) -> ModuleBase: + """Load an instance of a Caikit Model + + Args: + model_path (str): The model path to load from + model_id (str): The model's id + model_type (str): The type of model being load + finder (Optional[Union[str, ModelFinderBase]], optional): The ModelFinder to use for + loading. Defaults to None. + initializer (Optional[Union[str, ModelInitializerBase]], optional): The + ModelInitializer to use for loading. Defaults to None. + + Returns: + ModuleBase: a loaded model + """ + def load_model( self, model_id: str, @@ -91,34 +115,27 @@ def load_model( args = (local_model_path, model_id, model_type, finder, initializer) log.debug2("Loading model %s async", model_id) future_factory = partial( - self._load_thread_pool.submit, self._load_module, *args + self._load_thread_pool.submit, self._wrapped_load_model, *args ) model_builder.model_future_factory(future_factory) # Return the built model with the future handle return model_builder.build() - def _load_module( + def _wrapped_load_model( self, model_path: str, model_id: str, model_type: str, finder: Optional[Union[str, ModelFinderBase]] = None, initializer: Optional[Union[str, ModelInitializerBase]] = None, - ) -> LoadedModel: + ) -> Union[Batcher, ModuleBase]: try: log.info("", "Loading model '%s'", model_id) - # Only pass finder/initializer if they have values - load_kwargs = {} - if finder: - load_kwargs["finder"] = finder - if initializer: - load_kwargs["initializer"] = initializer - - # Load using the caikit.core - with CAIKIT_CORE_LOAD_DURATION_SUMMARY.labels(model_type=model_type).time(): - model = MODEL_MANAGER.load(model_path, **load_kwargs) + model = self.load_module_instance( + model_path, model_id, model_type, finder, initializer + ) # If this model needs batching, configure a Batcher to wrap it model = self._wrap_in_batcher_if_configured( @@ -126,6 +143,14 @@ def _load_module( model_type, model_id, ) + except CaikitRuntimeException as cre: + log_dict = { + "log_code": "", + "message": f"load failed to load model: {model_path} with error: {repr(cre)}", + "model_id": model_id, + } + log.error(log_dict) + raise cre except FileNotFoundError as fnfe: log_dict = { "log_code": "", @@ -165,13 +190,6 @@ def _load_module( return model - @classmethod - def get_instance(cls) -> "ModelLoader": - """This method returns the instance of Model Manager""" - if not cls.__instance: - cls.__instance = ModelLoader() - return cls.__instance - def _wrap_in_batcher_if_configured( self, caikit_core_model: ModuleBase, diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py index 3a8893447..8ff955956 100644 --- a/caikit/runtime/model_management/model_manager.py +++ b/caikit/runtime/model_management/model_manager.py @@ -35,9 +35,18 @@ from caikit.core import ModuleBase from caikit.core.exceptions import error_handler from caikit.core.model_management import ModelFinderBase, ModelInitializerBase +from caikit.runtime.model_management.factories import ( + model_loader_factory, + model_sizer_factory, +) from caikit.runtime.model_management.loaded_model import LoadedModel -from caikit.runtime.model_management.model_loader import ModelLoader -from caikit.runtime.model_management.model_sizer import ModelSizer +from caikit.runtime.model_management.model_loader_base import ModelLoaderBase +from caikit.runtime.model_management.model_sizer_base import ModelSizerBase +from caikit.runtime.names import ( + DEFAULT_LOADER_NAME, + DEFAULT_SIZER_NAME, + LOCAL_MODEL_TYPE, +) from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException log = alog.use_channel("MODEL-MANAGR") @@ -61,7 +70,6 @@ "Summary of the duration (in seconds) of loadModel RPCs", ["model_type"], ) -LOCAL_MODEL_TYPE = "LOCAL" class ModelManager: # pylint: disable=too-many-instance-attributes @@ -73,8 +81,6 @@ class ModelManager: # pylint: disable=too-many-instance-attributes __model_size_gauge_lock = threading.Lock() - _LOCAL_MODEL_TYPE = "standalone-model" - ## Construction ## @classmethod @@ -91,8 +97,30 @@ def __init__(self): ModelManager.__instance = self # Pull in a ModelLoader and ModelSizer - self.model_loader = ModelLoader.get_instance() - self.model_sizer = ModelSizer.get_instance() + loader_config = get_config().model_management.loaders.get( + DEFAULT_LOADER_NAME, {} + ) + error.value_check( + "", + isinstance(loader_config, dict), + "Unknown {}: {}", + "loader", + DEFAULT_LOADER_NAME, + ) + self.model_loader: ModelLoaderBase = model_loader_factory.construct( + loader_config, DEFAULT_LOADER_NAME + ) + sizer_config = get_config().model_management.sizers.get(DEFAULT_LOADER_NAME, {}) + error.value_check( + "", + isinstance(sizer_config, dict), + "Unknown {}: {}", + "sizer", + DEFAULT_SIZER_NAME, + ) + self.model_sizer: ModelSizerBase = model_sizer_factory.construct( + sizer_config, DEFAULT_LOADER_NAME + ) # In-memory mapping of model_id to LoadedModel instance self.loaded_models: Dict[str, LoadedModel] = {} @@ -433,7 +461,7 @@ def retrieve_model(self, model_id: str) -> ModuleBase: loaded_model = self.load_model( model_id=model_id, local_model_path=local_model_path, - model_type=self._LOCAL_MODEL_TYPE, + model_type=LOCAL_MODEL_TYPE, wait=True, retries=get_config().runtime.lazy_load_retries, ) @@ -510,7 +538,7 @@ def deploy_model( return self.load_model( model_id=model_id, local_model_path=model_dir, - model_type=self._LOCAL_MODEL_TYPE, + model_type=LOCAL_MODEL_TYPE, **kwargs, ) @@ -615,7 +643,7 @@ def _local_models_dir_sync(self, wait: bool = False, load: bool = True): self.load_model( model_id, model_path, - self._LOCAL_MODEL_TYPE, + LOCAL_MODEL_TYPE, wait=False, retries=get_config().runtime.lazy_load_retries, ) diff --git a/caikit/runtime/model_management/model_sizer.py b/caikit/runtime/model_management/model_sizer.py deleted file mode 100644 index 3f5b0c294..000000000 --- a/caikit/runtime/model_management/model_sizer.py +++ /dev/null @@ -1,121 +0,0 @@ -# Copyright The Caikit Authors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Standard -from pathlib import Path -from typing import Dict -import os - -# Third Party -import grpc - -# First Party -import alog - -# Local -from caikit import get_config -from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException - -log = alog.use_channel("MODEL-SIZER") - - -class ModelSizer: - """Model Loader class. The singleton class contains the core implementation details - for loading models in from S3.""" - - __instance = None - - def __init__(self): - # Re-instantiating this is a programming error - assert self.__class__.__instance is None, "This class is a singleton!" - ModelSizer.__instance = self - - # Cache of archive sizes: cos model path -> archive size in bytes - self._model_archive_sizes: Dict[str, int] = {} - - def get_model_size(self, model_id, local_model_path, model_type) -> int: - """ - Returns the estimated memory footprint of a model - Args: - model_id: The model identifier, used for informative logging - cos_model_path: The path to the model archive in S3 storage - model_type: The type of model, used to adjust the memory estimate - Returns: - The estimated size in bytes of memory that would be used by loading this model - """ - # Cache model's size - if local_model_path not in self._model_archive_sizes: - self._model_archive_sizes[local_model_path] = self.__get_archive_size( - model_id, local_model_path - ) - - return self.__estimate_with_multiplier( - model_id, model_type, self._model_archive_sizes[local_model_path] - ) - - def __estimate_with_multiplier(self, model_id, model_type, archive_size) -> int: - if ( - model_type - in get_config().inference_plugin.model_mesh.model_size_multipliers - ): - multiplier = ( - get_config().inference_plugin.model_mesh.model_size_multipliers[ - model_type - ] - ) - log.debug( - "Using size multiplier '%f' for model '%s' to estimate model size", - multiplier, - model_id, - ) - else: - multiplier = ( - get_config().inference_plugin.model_mesh.default_model_size_multiplier - ) - log.info( - "", - "No configured model size multiplier found for model type '%s' for model '%s'. " - "Using default multiplier '%f'", - model_type, - model_id, - multiplier, - ) - return int(archive_size * multiplier) - - def __get_archive_size(self, model_id, local_model_path) -> int: - try: - if os.path.isdir(local_model_path): - # Walk the directory to size all files - return sum( - file.stat().st_size - for file in Path(local_model_path).rglob("*") - if file.is_file() - ) - - # Probably just an archive file - return os.path.getsize(local_model_path) - except FileNotFoundError as ex: - message = ( - f"Failed to estimate size of model '{model_id}'," - f"file '{local_model_path}' not found" - ) - log.error("", message) - raise CaikitRuntimeException(grpc.StatusCode.NOT_FOUND, message) from ex - - @classmethod - def get_instance(cls) -> "ModelSizer": - """This method returns the instance of Model Manager""" - if not cls.__instance: - cls.__instance = ModelSizer() - return cls.__instance diff --git a/caikit/runtime/model_management/model_sizer_base.py b/caikit/runtime/model_management/model_sizer_base.py new file mode 100644 index 000000000..2931cf2dd --- /dev/null +++ b/caikit/runtime/model_management/model_sizer_base.py @@ -0,0 +1,42 @@ +# Copyright The Caikit Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Standard +import abc + +# First Party +import alog + +# Local +from caikit.core.toolkit.factory import FactoryConstructible + +log = alog.use_channel("MODEL-SIZER") + + +class ModelSizerBase(FactoryConstructible): + """Model Sizer Base class. This class contains the""" + + @abc.abstractmethod + def get_model_size( + self, model_id: str, local_model_path: str, model_type: str + ) -> int: + """ + Returns the estimated memory footprint of a model + Args: + model_id: The model identifier, used for informative logging + cos_model_path: The path to the model archive in S3 storage + model_type: The type of model, used to adjust the memory estimate + Returns: + The estimated size in bytes of memory that would be used by loading this model + """ diff --git a/caikit/runtime/names.py b/caikit/runtime/names.py index 3bdac38e1..71e333105 100644 --- a/caikit/runtime/names.py +++ b/caikit/runtime/names.py @@ -49,6 +49,11 @@ log = alog.use_channel("RNTM-NAMES") +################################# Model Management Names ####################### +LOCAL_MODEL_TYPE = "standalone-model" +DEFAULT_LOADER_NAME = "default" +DEFAULT_SIZER_NAME = "default" + ################################# Service Names ################################ diff --git a/tests/runtime/model_management/test_model_loader.py b/tests/runtime/model_management/test_model_loader.py index 7217fab6a..de6711ace 100644 --- a/tests/runtime/model_management/test_model_loader.py +++ b/tests/runtime/model_management/test_model_loader.py @@ -27,8 +27,10 @@ from caikit.core import ModuleConfig from caikit.core.module_backends import backend_types from caikit.core.modules import base, module +from caikit.core.toolkit.factory import FactoryConstructible from caikit.runtime.model_management.batcher import Batcher -from caikit.runtime.model_management.model_loader import ModelLoader +from caikit.runtime.model_management.core_model_loader import CoreModelLoader +from caikit.runtime.model_management.factories import model_loader_factory from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException from sample_lib.data_model import SampleInputType, SampleOutputType from sample_lib.modules.sample_task import SampleModule @@ -43,15 +45,19 @@ @contextmanager def temp_model_loader(): """Temporarily reset the ModelLoader singleton""" - real_singleton = ModelLoader.get_instance() - ModelLoader._ModelLoader__instance = None - yield ModelLoader.get_instance() - ModelLoader._ModelLoader__instance = real_singleton + yield construct_model_loader() @pytest.fixture def model_loader(): - return ModelLoader.get_instance() + yield construct_model_loader() + + +def construct_model_loader(): + model_loader: CoreModelLoader = model_loader_factory.construct( + get_config().model_management.loaders.default, "default" + ) + return model_loader def make_model_future(model_instance): @@ -161,10 +167,11 @@ def test_nonzip_extract_fails(model_loader): assert "config.yml" in context.value.message -def test_no_double_instantiation(): +def test_no_double_instantiation_of_thread_pools(): """Make sure trying to re-instantiate this singleton raises""" - with pytest.raises(Exception): - ModelLoader() + loader1 = construct_model_loader() + loader2 = construct_model_loader() + assert loader1._load_thread_pool is loader2._load_thread_pool def test_with_batching(model_loader): @@ -319,11 +326,11 @@ def test_load_model_succeed_after_retry(model_loader): """ failures = 2 fail_wrapper = TempFailWrapper( - model_loader._load_module, + model_loader.load_module_instance, num_failures=failures, exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Yikes!"), ) - with mock.patch.object(model_loader, "_load_module", fail_wrapper): + with mock.patch.object(model_loader, "load_module_instance", fail_wrapper): model_id = random_test_id() loaded_model = model_loader.load_model( model_id=model_id, @@ -342,11 +349,11 @@ def test_load_model_fail_callback_once(model_loader): """ failures = 3 fail_wrapper = TempFailWrapper( - model_loader._load_module, + model_loader.load_module_instance, num_failures=failures, exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Yikes!"), ) - with mock.patch.object(model_loader, "_load_module", fail_wrapper): + with mock.patch.object(model_loader, "load_module_instance", fail_wrapper): model_id = random_test_id() fail_cb = mock.MagicMock() loaded_model = model_loader.load_model( @@ -367,11 +374,13 @@ def test_load_model_loaded_status(model_loader): release_event = threading.Event() model_mock = mock.MagicMock() - def _load_module_mock(*_, **__): + def _load_module_instance_mock(*_, **__): release_event.wait() return model_mock - with mock.patch.object(model_loader, "_load_module", _load_module_mock): + with mock.patch.object( + model_loader, "load_module_instance", _load_module_instance_mock + ): loaded_model = model_loader.load_model( model_id=model_id, local_model_path=Fixtures.get_good_model_path(), diff --git a/tests/runtime/model_management/test_model_manager.py b/tests/runtime/model_management/test_model_manager.py index 847de5312..88e611eae 100644 --- a/tests/runtime/model_management/test_model_manager.py +++ b/tests/runtime/model_management/test_model_manager.py @@ -46,7 +46,7 @@ from tests.core.helpers import TestFinder from tests.fixtures import Fixtures from tests.runtime.conftest import deploy_good_model_files -import caikit.runtime.model_management.model_loader +import caikit.runtime.model_management.model_loader_base get_dynamic_module("caikit.core") ANY_MODEL_TYPE = "test-any-model-type" @@ -72,7 +72,7 @@ def temp_local_models_dir(workdir, model_manager=MODEL_MANAGER): def non_singleton_model_managers(num_mgrs=1, *args, **kwargs): with temp_config(*args, **kwargs): with patch( - "caikit.runtime.model_management.model_loader.MODEL_MANAGER", + "caikit.runtime.model_management.core_model_loader.MODEL_MANAGER", new_callable=CoreModelManager, ): instances = [] @@ -1269,13 +1269,13 @@ def test_periodic_sync_handles_temporary_errors(): ) as managers: manager = managers[0] flakey_loader = TempFailWrapper( - manager.model_loader._load_module, + manager.model_loader.load_module_instance, num_failures=1, exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Dang"), ) with patch.object( manager.model_loader, - "_load_module", + "load_module_instance", flakey_loader, ): assert manager._lazy_sync_timer is not None @@ -1307,13 +1307,13 @@ def test_lazy_load_handles_temporary_errors(): ) as managers: manager = managers[0] flakey_loader = TempFailWrapper( - manager.model_loader._load_module, + manager.model_loader.load_module_instance, num_failures=1, exc=CaikitRuntimeException(grpc.StatusCode.INTERNAL, "Dang"), ) with patch.object( manager.model_loader, - "_load_module", + "load_module_instance", flakey_loader, ): assert manager._lazy_sync_timer is None diff --git a/tests/runtime/model_management/test_model_sizer.py b/tests/runtime/model_management/test_model_sizer.py index 4ed891f43..ca266fe1b 100644 --- a/tests/runtime/model_management/test_model_sizer.py +++ b/tests/runtime/model_management/test_model_sizer.py @@ -22,7 +22,8 @@ # Local from caikit import get_config -from caikit.runtime.model_management.model_sizer import ModelSizer +from caikit.runtime.model_management.factories import model_sizer_factory +from caikit.runtime.model_management.model_sizer_base import ModelSizerBase from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException from tests.conftest import random_test_id, temp_config from tests.fixtures import Fixtures @@ -37,7 +38,9 @@ class TestModelSizer(unittest.TestCase): def setUp(self): """This method runs before each test begins to run""" - self.model_sizer = ModelSizer.get_instance() + self.model_sizer = model_sizer_factory.construct( + get_config().model_management.sizers.default, "default" + ) @staticmethod def _add_file(path, charsize) -> int: