feat(nodes): make delete on startup configurable for obj serializer

- The default is to not delete on startup - feels safer.
- The two services using this class _do_ delete on startup.
- The class has "ephemeral" removed from its name.
- Tests & app updated for this change.
This commit is contained in:
psychedelicious
2024-02-08 16:09:59 +11:00
parent ed772a7107
commit c58f8c3269
3 changed files with 67 additions and 27 deletions

View File

@ -5,7 +5,7 @@ from logging import Logger
import torch
from invokeai.app.services.item_storage.item_storage_memory import ItemStorageMemory
from invokeai.app.services.object_serializer.object_serializer_ephemeral_disk import ObjectSerializerEphemeralDisk
from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
from invokeai.app.services.shared.sqlite.sqlite_util import init_db
from invokeai.backend.model_manager.metadata import ModelMetadataStore
@ -90,9 +90,11 @@ class ApiDependencies:
image_records = SqliteImageRecordStorage(db=db)
images = ImageService()
invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size)
tensors = ObjectSerializerForwardCache(ObjectSerializerEphemeralDisk[torch.Tensor](output_folder / "tensors"))
tensors = ObjectSerializerForwardCache(
ObjectSerializerDisk[torch.Tensor](output_folder / "tensors", delete_on_startup=True)
)
conditioning = ObjectSerializerForwardCache(
ObjectSerializerEphemeralDisk[ConditioningFieldData](output_folder / "conditioning")
ObjectSerializerDisk[ConditioningFieldData](output_folder / "conditioning", delete_on_startup=True)
)
model_manager = ModelManagerService(config, logger)
model_record_service = ModelRecordServiceSQL(db=db)

View File

@ -22,26 +22,30 @@ class DeleteAllResult:
freed_space_bytes: float
class ObjectSerializerEphemeralDisk(ObjectSerializerBase[T]):
"""Provides a disk-backed ephemeral storage for arbitrary python objects. The storage is cleared at startup.
class ObjectSerializerDisk(ObjectSerializerBase[T]):
"""Provides a disk-backed storage for arbitrary python objects.
:param output_folder: The folder where the objects will be stored
:param delete_on_startup: If True, all objects in the output folder will be deleted on startup
"""
def __init__(self, output_dir: Path):
def __init__(self, output_dir: Path, delete_on_startup: bool = False):
super().__init__()
self._output_dir = output_dir
self._output_dir.mkdir(parents=True, exist_ok=True)
self._delete_on_startup = delete_on_startup
self.__obj_class_name: Optional[str] = None
def start(self, invoker: "Invoker") -> None:
self._invoker = invoker
delete_all_result = self._delete_all()
if delete_all_result.deleted_count > 0:
freed_space_in_mb = round(delete_all_result.freed_space_bytes / 1024 / 1024, 2)
self._invoker.services.logger.info(
f"Deleted {delete_all_result.deleted_count} {self._obj_class_name} files (freed {freed_space_in_mb}MB)"
)
if self._delete_on_startup:
delete_all_result = self._delete_all()
if delete_all_result.deleted_count > 0:
freed_space_in_mb = round(delete_all_result.freed_space_bytes / 1024 / 1024, 2)
self._invoker.services.logger.info(
f"Deleted {delete_all_result.deleted_count} {self._obj_class_name} files (freed {freed_space_in_mb}MB)"
)
def load(self, name: str) -> T:
file_path = self._get_path(name)