Model Cache API Flaw

[RyanJDick]

This issue is accurate as of v5.5.0.

There is a longstanding flaw in the design of the model cache API. It is the source of several issues and design quirks in the ModelCache, so I am taking the time to document it here so that it can be linked from the source code and referenced by devs.

This issue became particularly evident during the development of the partial model loading feature.

Typical Usage

Typical usage of the model cache API looks like this:

class SampleInvocation(BaseInvocation):
    model: ModelField = InputField(...)

    def invoke(self, context: InvocationContext):
        # context.models.load(...) does the following:
        # 1. Looks up the model info from the DB.
        # 2. Loads the model onto the CPU, using the ModelCache.
        # 3. Retuns a `LoadedModel` containing both the model metadata and the CPU-loaded model.
        model_info = context.models.load(self.model.model_id)

        # A context manager is used to load the model into VRAM using the
        # model cache and 'lock' it there for the lifetime of the context
        # manager.
        with model_info.model_on_device() as (cached_weights, model):
        # Or using the deprecated API:
        # with model_info as model:
            y = model.forward(...)

Flaw

The main problem is that after calling context.models.load(...) the ModelCache has no way of knowing whether:

1. this model needs to be kept in the RAM cache, because it will soon be loaded into VRAM, or
2. this model is only used on the CPU and can be freed from the RAM cache to make room for other models

As a result of this flaw, the following reasonable-looking code could fail:

class SampleInvocation(BaseInvocation):
    model_1: ModelField = InputField(...)
    model_2: ModelField = InputField(...)

    def invoke(self, context: InvocationContext):
	    model_1_info = context.models.load(self.model_1.model_id)
	    # Loading model_2 into the CPU could cause model_1 to be ejected
	    # from the RAM cache to make room for model_2.
	    model_2_info = context.models.load(self.model_2.model_id)

		# If model_1 was ejected from the RAM cache, then attempting to
		# load it into VRAM will fail.
		with (model_1_info as model_1, model_2_info as model_2):
			# Do something with model_1 and model_2...
			...

If this flaw is not handled carefully, the following error will occur:

[2025-01-03 10:25:25,390]::[InvokeAI]::ERROR --> Error while invoking session 93e35ea6-4ed8-4535-b7ff-f4d6d3918afd, invocation 36faba26-ce91-483f-b7fe-59bcda83e07f (compel): 'a67e6e91-6b2a-4d38-b69a-db8c27ee98a8:tokenizer'
[2025-01-03 10:25:25,390]::[InvokeAI]::ERROR --> Traceback (most recent call last):
  File "/home/ryan/src/InvokeAI/invokeai/app/services/session_processor/session_processor_default.py", line 129, in run_node
    output = invocation.invoke_internal(context=context, services=self._services)
  File "/home/ryan/src/InvokeAI/invokeai/app/invocations/baseinvocation.py", line 300, in invoke_internal
    output = self.invoke(context)
  File "/home/ryan/.pyenv/versions/3.10.12/envs/InvokeAI_3.10.12/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/home/ryan/src/InvokeAI/invokeai/app/invocations/compel.py", line 81, in invoke
    with (
  File "/home/ryan/src/InvokeAI/invokeai/backend/model_manager/load/load_base.py", line 60, in __enter__
    self._cache.lock(self._cache_record.key, None)
  File "/home/ryan/src/InvokeAI/invokeai/backend/model_manager/load/model_cache/model_cache.py", line 205, in lock
    cache_entry = self._cached_models[key]
KeyError: 'a67e6e91-6b2a-4d38-b69a-db8c27ee98a8:tokenizer'

Current Workaround

The current workaround for this issue is to allow models to be locked in VRAM even if they have been dropped from the RAM cache.

While this works, it is problematic in the context of partial model loading, because the model cache does not have the ability to partially-unload this model if it needs more VRAM.

Other Workarounds

Increase ram cache size

Increasing the ram config in invokeai.yaml increases the size of the RAM cache and decreases the probability that a model will get ejected between loading to RAM and loading to VRAM.

Avoid the offending model access pattern

Of course, avoiding the offending model access pattern in invocation code will prevent this issue from being encountered.

In practice, it is sometimes inconvenient and less performant to do this. I.e. it may require accessing the model info from the DB multiple times or loading the model into the CPU multiple times.

Reference Counting

In the past, we tried to workaround this flaw in the API design by doing manual reference counting on the CPU-loaded models to determine which could be safely dropped from the RAM cache. This approach was very error-prone and the source of several bugs. It also placed some burden on the invocation author to understand how this reference counting was happening under the hood if they wanted to achieve optimal memory behaviour.

Handle in ModelCache.lock() or LoadedModelWithoutConfig

It is tempting to workaround this issue by catching the KeyError and re-inserting the model into the cache (since we still have a reference to it). This feels risky. You would then have multiple CacheRecords floating around that reference the same underlying model.

Solution

A proper solution to this problem should achieve the following:

• Separate the model info DB lookup from model loading
• Enable the user to 'lock' the model on any device.