Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(nodes): allow _delete_all in obj serializer to be called at any time

Browse Source 
`_delete_all` logged how many items it deleted, and had to be called _after_ service start bc it needed access to logger.

Move the logger call to the startup method and return the the deleted stats from `_delete_all`. This lets `_delete_all` be called at any time.
This commit is contained in:
psychedelicious 2024-02-08 00:18:58 +11:00
parent 507aeac8a5
commit 23de78ec9f
1 changed files with 15 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
invokeai/app/services/object_serializer/object_serializer_ephemeral_disk.py
View File

@ -1,4 +1,5 @@
import typing import typing
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Optional, TypeVar from typing import Optional, TypeVar
@ -12,6 +13,12 @@ from invokeai.app.util.misc import uuid_string
T = TypeVar("T") T = TypeVar("T")
@dataclass
class DeleteAllResult:
deleted_count: int
freed_space_bytes: float
class ObjectSerializerEphemeralDisk(ObjectSerializerBase[T]): class ObjectSerializerEphemeralDisk(ObjectSerializerBase[T]):
"""Provides a disk-backed ephemeral storage for arbitrary python objects. The storage is cleared at startup. """Provides a disk-backed ephemeral storage for arbitrary python objects. The storage is cleared at startup.
@ -26,7 +33,12 @@ class ObjectSerializerEphemeralDisk(ObjectSerializerBase[T]):
def start(self, invoker: Invoker) -> None: def start(self, invoker: Invoker) -> None:
self._invoker = invoker self._invoker = invoker
self._delete_all() 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: def load(self, name: str) -> T:
file_path = self._get_path(name) file_path = self._get_path(name)
@ -58,18 +70,14 @@ class ObjectSerializerEphemeralDisk(ObjectSerializerBase[T]):
def _new_name(self) -> str: def _new_name(self) -> str:
return f"{self._obj_class_name}_{uuid_string()}" return f"{self._obj_class_name}_{uuid_string()}"
def _delete_all(self) -> None: def _delete_all(self) -> DeleteAllResult:
""" """
Deletes all objects from disk. Deletes all objects from disk.
Must be called after we have access to `self._invoker` (e.g. in `start()`).
""" """
# We could try using a temporary directory here, but they aren't cleared in the event of a crash, so we'd have # We could try using a temporary directory here, but they aren't cleared in the event of a crash, so we'd have
# to manually clear them on startup anyways. This is a bit simpler and more reliable. # to manually clear them on startup anyways. This is a bit simpler and more reliable.
if not self._invoker:
raise ValueError("Invoker is not set. Must call `start()` first.")
deleted_count = 0 deleted_count = 0
freed_space = 0 freed_space = 0
for file in Path(self._output_dir).glob("*"): for file in Path(self._output_dir).glob("*"):
@ -77,8 +85,4 @@ class ObjectSerializerEphemeralDisk(ObjectSerializerBase[T]):
freed_space += file.stat().st_size freed_space += file.stat().st_size
deleted_count += 1 deleted_count += 1
file.unlink() file.unlink()
if deleted_count > 0: return DeleteAllResult(deleted_count, freed_space)
freed_space_in_mb = round(freed_space / 1024 / 1024, 2)
self._invoker.services.logger.info(
f"Deleted {deleted_count} {self._obj_class_name} files (freed {freed_space_in_mb}MB)"
)