diff --git a/src/neptune/new/internal/operation_processors/operation_storage.py b/src/neptune/new/internal/operation_processors/operation_storage.py index bc2726a0e..39eb0293d 100644 --- a/src/neptune/new/internal/operation_processors/operation_storage.py +++ b/src/neptune/new/internal/operation_processors/operation_storage.py @@ -18,11 +18,13 @@ ] import os +import shutil from pathlib import Path from neptune.new.constants import NEPTUNE_DATA_DIRECTORY from neptune.new.internal.container_type import ContainerType from neptune.new.internal.id_formats import UniqueId +from neptune.new.internal.utils.logger import logger class OperationStorage: @@ -44,3 +46,15 @@ def upload_path(self) -> Path: @staticmethod def _get_container_dir(type_dir: str, container_id: UniqueId, container_type: ContainerType): return f"{NEPTUNE_DATA_DIRECTORY}/{type_dir}/{container_type.create_dir_name(container_id)}" + + def close(self): + shutil.rmtree(self.data_path, ignore_errors=True) + + parent = self.data_path.parent + + files = os.listdir(parent) + if len(files) == 0: + try: + os.rmdir(parent) + except OSError: + logger.debug(f"Cannot remove directory: {parent}") diff --git a/src/neptune/new/internal/operation_processors/sync_operation_processor.py b/src/neptune/new/internal/operation_processors/sync_operation_processor.py index 31d83cfbe..68536f0fb 100644 --- a/src/neptune/new/internal/operation_processors/sync_operation_processor.py +++ b/src/neptune/new/internal/operation_processors/sync_operation_processor.py @@ -59,4 +59,4 @@ def start(self): pass def stop(self, seconds: Optional[float] = None): - pass + self._operation_storage.close() diff --git a/tests/neptune/new/attributes/atoms/test_file.py b/tests/neptune/new/attributes/atoms/test_file.py index 56984192b..7f46648f3 100644 --- a/tests/neptune/new/attributes/atoms/test_file.py +++ b/tests/neptune/new/attributes/atoms/test_file.py @@ -161,3 +161,13 @@ def test_fetch_extension(self): var = File(exp, path) var.assign(value, wait=wait) self.assertEqual(expected_ext, var.fetch_extension()) + + def test_clean_files_on_close(self): + run = self._create_run() + data_path = run._op_processor._operation_storage.data_path + + assert os.path.exists(data_path) + + run.stop() + + assert not os.path.exists(data_path)