feat(backend): display freed space when cleaning DB

This commit is contained in:
psychedelicious 2023-11-30 20:07:53 +11:00 committed by Kent Keirsey
parent 1fd6666682
commit 0228aba06f
2 changed files with 12 additions and 11 deletions

View File

@ -33,7 +33,7 @@ class DiskLatentsStorage(LatentsStorageBase):
if deleted_latents_count > 0:
freed_space_in_mb = round(freed_space / 1024 / 1024, 2)
self._invoker.services.logger.info(
f"Deleted {deleted_latents_count} latents files, freeing {freed_space_in_mb}MB"
f"Deleted {deleted_latents_count} latents files (freed {freed_space_in_mb}MB)"
)
def get(self, name: str) -> torch.Tensor:

View File

@ -1,6 +1,7 @@
import sqlite3
import threading
from logging import Logger
from pathlib import Path
from invokeai.app.services.config import InvokeAIAppConfig
@ -8,25 +9,20 @@ sqlite_memory = ":memory:"
class SqliteDatabase:
conn: sqlite3.Connection
lock: threading.RLock
_logger: Logger
_config: InvokeAIAppConfig
def __init__(self, config: InvokeAIAppConfig, logger: Logger):
self._logger = logger
self._config = config
if self._config.use_memory_db:
location = sqlite_memory
self.db_path = sqlite_memory
logger.info("Using in-memory database")
else:
db_path = self._config.db_path
db_path.parent.mkdir(parents=True, exist_ok=True)
location = str(db_path)
self._logger.info(f"Using database at {location}")
self.db_path = str(db_path)
self._logger.info(f"Using database at {self.db_path}")
self.conn = sqlite3.connect(location, check_same_thread=False)
self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
self.lock = threading.RLock()
self.conn.row_factory = sqlite3.Row
@ -37,10 +33,15 @@ class SqliteDatabase:
def clean(self) -> None:
try:
if self.db_path == sqlite_memory:
return
initial_db_size = Path(self.db_path).stat().st_size
self.lock.acquire()
self.conn.execute("VACUUM;")
self.conn.commit()
self._logger.info("Cleaned database")
final_db_size = Path(self.db_path).stat().st_size
freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2)
self._logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)")
except Exception as e:
self._logger.error(f"Error cleaning database: {e}")
raise e