From dc75bb8b7c406740908db1b7f9a68df72f6bc943 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Fri, 19 Apr 2024 10:42:26 -0400 Subject: [PATCH 1/6] add an initialization function that removes dangling tmpdirs from outputs/tensors --- invokeai/app/api/dependencies.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index 9a6c7416f69..0bbe524506e 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -1,6 +1,8 @@ # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) +import shutil from logging import Logger +from pathlib import Path import torch @@ -52,6 +54,14 @@ def check_internet() -> bool: return False +def cleanup_tmpdirs(parent_folder: Path) -> None: + # Remove dangling tempdirs that might have been left over + # from an earlier unplanned shutdown. + for d in parent_folder.glob("tmp*"): + if d.is_dir(): + shutil.rmtree(d) + + logger = InvokeAILogger.get_logger() @@ -70,7 +80,7 @@ def initialize(config: InvokeAIAppConfig, event_handler_id: int, logger: Logger raise ValueError("Output folder is not set") image_files = DiskImageFileStorage(f"{output_folder}/images") - + tensor_folder = output_folder / "tensors" model_images_folder = config.models_path db = init_db(config=config, logger=logger, image_files=image_files) @@ -78,6 +88,7 @@ def initialize(config: InvokeAIAppConfig, event_handler_id: int, logger: Logger configuration = config logger = logger + cleanup_tmpdirs(tensor_folder) board_image_records = SqliteBoardImageRecordStorage(db=db) board_images = BoardImagesService() board_records = SqliteBoardRecordStorage(db=db) @@ -87,9 +98,7 @@ def initialize(config: InvokeAIAppConfig, event_handler_id: int, logger: Logger image_records = SqliteImageRecordStorage(db=db) images = ImageService() invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size) - tensors = ObjectSerializerForwardCache( - ObjectSerializerDisk[torch.Tensor](output_folder / "tensors", ephemeral=True) - ) + tensors = ObjectSerializerForwardCache(ObjectSerializerDisk[torch.Tensor](tensor_folder, ephemeral=True)) conditioning = ObjectSerializerForwardCache( ObjectSerializerDisk[ConditioningFieldData](output_folder / "conditioning", ephemeral=True) ) From 81e7a2755253dcd32ce41651614ae7c83042eebc Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 22 Apr 2024 21:33:17 -0400 Subject: [PATCH 2/6] moved cleanup routine into object_serializer_disk.py --- invokeai/app/api/dependencies.py | 11 ----------- .../object_serializer/object_serializer_disk.py | 12 ++++++++++++ 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index 0bbe524506e..346b2742484 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -1,8 +1,6 @@ # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) -import shutil from logging import Logger -from pathlib import Path import torch @@ -54,14 +52,6 @@ def check_internet() -> bool: return False -def cleanup_tmpdirs(parent_folder: Path) -> None: - # Remove dangling tempdirs that might have been left over - # from an earlier unplanned shutdown. - for d in parent_folder.glob("tmp*"): - if d.is_dir(): - shutil.rmtree(d) - - logger = InvokeAILogger.get_logger() @@ -88,7 +78,6 @@ def initialize(config: InvokeAIAppConfig, event_handler_id: int, logger: Logger configuration = config logger = logger - cleanup_tmpdirs(tensor_folder) board_image_records = SqliteBoardImageRecordStorage(db=db) board_images = BoardImagesService() board_records = SqliteBoardRecordStorage(db=db) diff --git a/invokeai/app/services/object_serializer/object_serializer_disk.py b/invokeai/app/services/object_serializer/object_serializer_disk.py index 935fec30605..16e5ffa9c97 100644 --- a/invokeai/app/services/object_serializer/object_serializer_disk.py +++ b/invokeai/app/services/object_serializer/object_serializer_disk.py @@ -1,3 +1,4 @@ +import shutil import tempfile import typing from dataclasses import dataclass @@ -81,5 +82,16 @@ def __del__(self) -> None: # In case the service is not properly stopped, clean up the temporary directory when the class instance is GC'd. self._tempdir_cleanup() + @classmethod + def _cleanup_dangling_temporary_dirs(cls, directory: Path): + # Remove dangling tempdirs that might have been left over + # from an earlier unplanned shutdown. + for d in directory.glob("tmp*"): + if d.is_dir(): + shutil.rmtree(d) + + def start(self, invoker: "Invoker") -> None: + self._cleanup_dangling_temporary_dirs(self._base_output_dir) + def stop(self, invoker: "Invoker") -> None: self._tempdir_cleanup() From 1e031ca4dd0d195fd0439f422becd11c3746c30e Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 23 Apr 2024 16:54:28 +1000 Subject: [PATCH 3/6] tidy(api): reverted unnecessary changes in dependencies.py --- invokeai/app/api/dependencies.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index 346b2742484..9a6c7416f69 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -70,7 +70,7 @@ def initialize(config: InvokeAIAppConfig, event_handler_id: int, logger: Logger raise ValueError("Output folder is not set") image_files = DiskImageFileStorage(f"{output_folder}/images") - tensor_folder = output_folder / "tensors" + model_images_folder = config.models_path db = init_db(config=config, logger=logger, image_files=image_files) @@ -87,7 +87,9 @@ def initialize(config: InvokeAIAppConfig, event_handler_id: int, logger: Logger image_records = SqliteImageRecordStorage(db=db) images = ImageService() invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size) - tensors = ObjectSerializerForwardCache(ObjectSerializerDisk[torch.Tensor](tensor_folder, ephemeral=True)) + tensors = ObjectSerializerForwardCache( + ObjectSerializerDisk[torch.Tensor](output_folder / "tensors", ephemeral=True) + ) conditioning = ObjectSerializerForwardCache( ObjectSerializerDisk[ConditioningFieldData](output_folder / "conditioning", ephemeral=True) ) From d0811a730f0da40a68f74093ff1a360aa28ad924 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 23 Apr 2024 16:55:46 +1000 Subject: [PATCH 4/6] fix(app): only clear tempdirs if ephemeral and before creating tempdir Also, this needs to happen in init, else it deletes the temp dir created in init --- .../object_serializer/object_serializer_disk.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/invokeai/app/services/object_serializer/object_serializer_disk.py b/invokeai/app/services/object_serializer/object_serializer_disk.py index 16e5ffa9c97..df3327d20f3 100644 --- a/invokeai/app/services/object_serializer/object_serializer_disk.py +++ b/invokeai/app/services/object_serializer/object_serializer_disk.py @@ -36,6 +36,12 @@ def __init__(self, output_dir: Path, ephemeral: bool = False): self._ephemeral = ephemeral self._base_output_dir = output_dir self._base_output_dir.mkdir(parents=True, exist_ok=True) + + if self._ephemeral: + # Remove dangling tempdirs that might have been left over from an earlier unplanned shutdown. + for temp_dir in filter(Path.is_dir, self._base_output_dir.glob("tmp*")): + shutil.rmtree(temp_dir) + # Must specify `ignore_cleanup_errors` to avoid fatal errors during cleanup on Windows self._tempdir = ( tempfile.TemporaryDirectory(dir=self._base_output_dir, ignore_cleanup_errors=True) if ephemeral else None @@ -82,16 +88,5 @@ def __del__(self) -> None: # In case the service is not properly stopped, clean up the temporary directory when the class instance is GC'd. self._tempdir_cleanup() - @classmethod - def _cleanup_dangling_temporary_dirs(cls, directory: Path): - # Remove dangling tempdirs that might have been left over - # from an earlier unplanned shutdown. - for d in directory.glob("tmp*"): - if d.is_dir(): - shutil.rmtree(d) - - def start(self, invoker: "Invoker") -> None: - self._cleanup_dangling_temporary_dirs(self._base_output_dir) - def stop(self, invoker: "Invoker") -> None: self._tempdir_cleanup() From a66d34d3d64d33beaa4f4f760efd9dfadcb124d8 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 23 Apr 2024 16:56:00 +1000 Subject: [PATCH 5/6] tidy(app): remove unused class --- .../services/object_serializer/object_serializer_disk.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/invokeai/app/services/object_serializer/object_serializer_disk.py b/invokeai/app/services/object_serializer/object_serializer_disk.py index df3327d20f3..8edd29e1505 100644 --- a/invokeai/app/services/object_serializer/object_serializer_disk.py +++ b/invokeai/app/services/object_serializer/object_serializer_disk.py @@ -1,7 +1,6 @@ import shutil import tempfile import typing -from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Optional, TypeVar @@ -18,12 +17,6 @@ T = TypeVar("T") -@dataclass -class DeleteAllResult: - deleted_count: int - freed_space_bytes: float - - class ObjectSerializerDisk(ObjectSerializerBase[T]): """Disk-backed storage for arbitrary python objects. Serialization is handled by `torch.save` and `torch.load`. From 8f8514dd680149aeb132e55d9eae4b47cf6501d6 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 23 Apr 2024 17:00:35 +1000 Subject: [PATCH 6/6] tests: add object serializer test for dangling folders - Ensure they are deleted on init if ephemeral - Ensure they are _not_ deleted on init if _not_ ephemeral --- tests/test_object_serializer_disk.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_object_serializer_disk.py b/tests/test_object_serializer_disk.py index 125534c5002..84c6e876fcb 100644 --- a/tests/test_object_serializer_disk.py +++ b/tests/test_object_serializer_disk.py @@ -99,6 +99,20 @@ def test_obj_serializer_ephemeral_writes_to_tempdir(tmp_path: Path): assert not Path(tmp_path, obj_1_name).exists() +def test_obj_serializer_ephemeral_deletes_dangling_tempdirs_on_init(tmp_path: Path): + tempdir = tmp_path / "tmpdir" + tempdir.mkdir() + ObjectSerializerDisk[MockDataclass](tmp_path, ephemeral=True) + assert not tempdir.exists() + + +def test_obj_serializer_does_not_delete_tempdirs_on_init(tmp_path: Path): + tempdir = tmp_path / "tmpdir" + tempdir.mkdir() + ObjectSerializerDisk[MockDataclass](tmp_path, ephemeral=False) + assert tempdir.exists() + + def test_obj_serializer_disk_different_types(tmp_path: Path): obj_serializer_1 = ObjectSerializerDisk[MockDataclass](tmp_path) obj_1 = MockDataclass(foo="bar")