flag to disable deletion of local parent folder

CHANGELOG.md

@@ -1,5 +1,10 @@
 ## [UNRELEASED] neptune 1.8.3
 
+### Changes
+- Safety (errors suppressing) execution mode ([#1503])(https://github.com/neptune-ai/neptune-client/pull/1503)
+- Allow to disable handling of remote signals ([#1508])(https://github.com/neptune-ai/neptune-client/pull/1508)
+- Allow to disable deletion of local parent folder ([#1511])(https://github.com/neptune-ai/neptune-client/pull/1511)
+
 ### Fixes
 - Added more safe checking to last ack ([#1510](https://github.com/neptune-ai/neptune-client/pull/1510))
 

src/neptune/envs.py

@@ -34,6 +34,7 @@
     "NEPTUNE_NON_RAISING_ON_DISK_ISSUE",
     "NEPTUNE_ENABLE_DEFAULT_ASYNC_LAG_CALLBACK",
     "NEPTUNE_ENABLE_DEFAULT_ASYNC_NO_PROGRESS_CALLBACK",
+    "NEPTUNE_DISABLE_PARENT_DIR_DELETION",
 ]
 
 from neptune.common.envs import (
@@ -77,4 +78,6 @@
 
 NEPTUNE_NON_RAISING_ON_DISK_ISSUE = "NEPTUNE_NON_RAISING_ON_DISK_ISSUE"
 
+NEPTUNE_DISABLE_PARENT_DIR_DELETION = "NEPTUNE_DISABLE_PARENT_DIR_DELETION"
+
 S3_ENDPOINT_URL = "S3_ENDPOINT_URL"

src/neptune/internal/disk_queue.py

@@ -33,6 +33,7 @@
 )
 
 from neptune.exceptions import MalformedOperation
+from neptune.internal.utils.files import remove_parent_folder_if_allowed
 from neptune.internal.utils.json_file_splitter import JsonFileSplitter
 from neptune.internal.utils.sync_offset_file import SyncOffsetFile
 
@@ -183,19 +184,7 @@ def cleanup_if_empty(self) -> None:
     def _remove_data(self):
         path = self._dir_path
         shutil.rmtree(path, ignore_errors=True)
-
-        parent = path.parent
-
-        try:
-            files = os.listdir(parent)
-        except FileNotFoundError:
-            files = []
-
-        if len(files) == 0:
-            try:
-                os.rmdir(parent)
-            except OSError:
-                _logger.info(f"Cannot remove directory: {parent}")
+        remove_parent_folder_if_allowed(path)
 
     def wait_for_empty(self, seconds: Optional[float] = None) -> bool:
         with self._empty_cond:

src/neptune/internal/operation_processors/operation_storage.py

@@ -23,7 +23,7 @@
 from neptune.constants import NEPTUNE_DATA_DIRECTORY
 from neptune.internal.container_type import ContainerType
 from neptune.internal.id_formats import UniqueId
-from neptune.internal.utils.logger import logger
+from neptune.internal.utils.files import remove_parent_folder_if_allowed
 
 
 def get_container_dir(
@@ -56,12 +56,4 @@ def upload_path(self) -> Path:
 
     def cleanup(self) -> None:
         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}")
+        remove_parent_folder_if_allowed(self.data_path)

src/neptune/internal/utils/files.py

@@ -0,0 +1,40 @@
+#
+# Copyright (c) 2023, Neptune Labs Sp. z o.o.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+__all__ = ["remove_parent_folder_if_allowed"]
+
+import os
+from pathlib import Path
+
+from neptune.envs import NEPTUNE_DISABLE_PARENT_DIR_DELETION
+from neptune.internal.utils.logger import logger
+
+
+def remove_parent_folder_if_allowed(path: Path) -> None:
+    deletion_disabled = os.getenv(NEPTUNE_DISABLE_PARENT_DIR_DELETION, "false").lower() in ("true", "1", "t")
+
+    if not deletion_disabled:
+        parent = path.parent
+
+        try:
+            files = os.listdir(parent)
+        except FileNotFoundError:
+            files = []
+
+        if len(files) == 0:
+            try:
+                os.rmdir(parent)
+            except OSError:
+                logger.debug(f"Cannot remove directory: {parent}")

tests/unit/neptune/new/attributes/atoms/test_file.py

@@ -36,6 +36,7 @@
     FileSetVal,
 )
 from neptune.common.utils import IS_WINDOWS
+from neptune.envs import NEPTUNE_DISABLE_PARENT_DIR_DELETION
 from neptune.internal.operation import (
     UploadFile,
     UploadFileSet,
@@ -186,4 +187,17 @@ def test_clean_files_on_close(self):
 
             run.stop()
 
-            assert not os.path.exists(data_path)
+            assert not os.path.exists(data_path)  # exec folder
+            assert not os.path.exists(data_path.parent)  # run folder
+
+    @patch.dict(os.environ, {NEPTUNE_DISABLE_PARENT_DIR_DELETION: "True"})
+    def test_clean_files_on_close_when_cleanup_disabled(self):
+        with self._exp() as 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)  # exec folder
+            assert os.path.exists(data_path.parent)  # run folder