Merge pull request #1519 from neptune-ai/at/more_defensive_full_disk

Make ensure_disk_not_full more reliable for env specific errors

CHANGELOG.md

@@ -4,6 +4,7 @@
 - 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))
+- Made the disk checking more reliable for env specific errors ([#1519](https://github.com/neptune-ai/neptune-client/pull/1519))
 
 ### Fixes
 - Added more safe checking to last ack ([#1510](https://github.com/neptune-ai/neptune-client/pull/1510))

src/neptune/internal/utils/disk_full.py

@@ -27,6 +27,7 @@
 )
 
 import psutil
+from psutil import Error
 
 from neptune.common.warnings import (
     NeptuneWarning,
@@ -70,7 +71,7 @@ def wrapper(*args: Tuple, **kwargs: Dict[str, Any]) -> None:
             try:
                 if max_disk_utilization:
                     current_utilization = get_disk_utilization_percent()
-                    if current_utilization > max_disk_utilization:
+                    if current_utilization >= max_disk_utilization:
                         warn_once(
                             f"Max disk utilization {max_disk_utilization}% exceeded with {current_utilization}."
                             f" Neptune will not be saving your data.",
@@ -79,7 +80,7 @@ def wrapper(*args: Tuple, **kwargs: Dict[str, Any]) -> None:
                         return
 
                 func(*args, **kwargs)
-            except IOError:
+            except (OSError, Error):
                 warn_once("Encountered disk issue and Neptune will not be saving your data.", exception=NeptuneWarning)
         else:
             return func(*args, **kwargs)

tests/unit/neptune/new/internal/utils/test_full_disk.py

@@ -0,0 +1,87 @@
+#
+# 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.
+#
+import os
+import unittest
+from io import UnsupportedOperation
+
+from mock import (
+    MagicMock,
+    patch,
+)
+from psutil import (
+    AccessDenied,
+    Error,
+)
+
+from neptune.envs import NEPTUNE_NON_RAISING_ON_DISK_ISSUE
+from neptune.internal.utils.disk_full import ensure_disk_not_full
+
+
+class TestDiskFull(unittest.TestCase):
+
+    # Catching OSError that's base error for all OS and IO errors. More info here: https://peps.python.org/pep-3151
+    # Additionally, catching specific psutil's base error - psutil.Error.
+    # More info about psutil.Error here: https://psutil.readthedocs.io/en/latest/index.html#psutil.Error
+    @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"})
+    @patch("neptune.internal.utils.disk_full.get_max_percentage_from_env")
+    @patch("neptune.internal.utils.disk_full.get_disk_utilization_percent")
+    def test_suppressing_of_env_errors(self, get_disk_utilization_percent, get_max_percentage_from_env):
+        get_max_percentage_from_env.return_value = 42
+
+        env_errors = [OSError(), IOError(), EnvironmentError(), UnsupportedOperation(), Error(), AccessDenied()]
+        for error in env_errors:
+            mocked_func = MagicMock()
+            wrapped_func = ensure_disk_not_full(mocked_func)
+            get_disk_utilization_percent.side_effect = error
+
+            wrapped_func()  # asserting is not required as expecting that any error will be caught
+            mocked_func.assert_not_called()
+
+        non_env_errors = [ValueError(), OverflowError()]
+        for error in non_env_errors:
+            mocked_func = MagicMock()
+            wrapped_func = ensure_disk_not_full(mocked_func)
+            get_disk_utilization_percent.side_effect = error
+
+            with self.assertRaises(BaseException):
+                wrapped_func()
+            mocked_func.assert_not_called()
+
+    @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"})
+    @patch("neptune.internal.utils.disk_full.get_max_percentage_from_env")
+    @patch("neptune.internal.utils.disk_full.get_disk_utilization_percent")
+    def test_not_called_with_usage_100_percent(self, get_disk_utilization_percent, get_max_percentage_from_env):
+        get_max_percentage_from_env.return_value = 100
+        get_disk_utilization_percent.return_value = 100
+        mocked_func = MagicMock()
+        wrapped_func = ensure_disk_not_full(mocked_func)
+
+        wrapped_func()
+
+        mocked_func.assert_not_called()
+
+    @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"})
+    @patch("neptune.internal.utils.disk_full.get_max_percentage_from_env")
+    @patch("neptune.internal.utils.disk_full.get_disk_utilization_percent")
+    def test_called_when_usage_less_than_limit(self, get_disk_utilization_percent, get_max_percentage_from_env):
+        get_max_percentage_from_env.return_value = 100
+        get_disk_utilization_percent.return_value = 99
+        mocked_func = MagicMock()
+        wrapped_func = ensure_disk_not_full(mocked_func)
+
+        wrapped_func()
+
+        mocked_func.assert_called_once()