From 6ebb77827ac3c3537dd6c50eeeff7ffbeaaa1c99 Mon Sep 17 00:00:00 2001 From: Artsiom Tserashkovich Date: Mon, 30 Oct 2023 10:51:51 +0100 Subject: [PATCH 1/6] Rework disk usage limitation: use single env var instead of 2 / add additional validations for input value / add additional warnings --- src/neptune/envs.py | 7 +- .../async_operation_processor.py | 4 +- .../offline_operation_processor.py | 4 +- .../sync_operation_processor.py | 4 +- src/neptune/internal/utils/disk_full.py | 88 --------------- .../internal/utils/disk_utilization.py | 102 ++++++++++++++++++ .../internal/utils/test_disk_utilization.py | 100 +++++++++++++++++ .../new/internal/utils/test_full_disk.py | 87 --------------- 8 files changed, 210 insertions(+), 186 deletions(-) delete mode 100644 src/neptune/internal/utils/disk_full.py create mode 100644 src/neptune/internal/utils/disk_utilization.py create mode 100644 tests/unit/neptune/new/internal/utils/test_disk_utilization.py delete mode 100644 tests/unit/neptune/new/internal/utils/test_full_disk.py diff --git a/src/neptune/envs.py b/src/neptune/envs.py index df9df83ff..c262b8413 100644 --- a/src/neptune/envs.py +++ b/src/neptune/envs.py @@ -31,8 +31,7 @@ "NEPTUNE_FETCH_TABLE_STEP_SIZE", "NEPTUNE_SYNC_AFTER_STOP_TIMEOUT", "NEPTUNE_REQUEST_TIMEOUT", - "NEPTUNE_MAX_DISK_UTILIZATION", - "NEPTUNE_NON_RAISING_ON_DISK_ISSUE", + "NEPTUNE_LIMIT_DISK_UTILIZATION", "NEPTUNE_ENABLE_DEFAULT_ASYNC_LAG_CALLBACK", "NEPTUNE_ENABLE_DEFAULT_ASYNC_NO_PROGRESS_CALLBACK", "NEPTUNE_DISABLE_PARENT_DIR_DELETION", @@ -78,9 +77,7 @@ NEPTUNE_ENABLE_DEFAULT_ASYNC_NO_PROGRESS_CALLBACK = "NEPTUNE_ENABLE_DEFAULT_ASYNC_NO_PROGRESS_CALLBACK" -NEPTUNE_MAX_DISK_UTILIZATION = "NEPTUNE_MAX_DISK_UTILIZATION" - -NEPTUNE_NON_RAISING_ON_DISK_ISSUE = "NEPTUNE_NON_RAISING_ON_DISK_ISSUE" +NEPTUNE_LIMIT_DISK_UTILIZATION = "NEPTUNE_LIMIT_DISK_UTILIZATION" NEPTUNE_DISABLE_PARENT_DIR_DELETION = "NEPTUNE_DISABLE_PARENT_DIR_DELETION" diff --git a/src/neptune/internal/operation_processors/async_operation_processor.py b/src/neptune/internal/operation_processors/async_operation_processor.py index cffed58e2..84d4c0054 100644 --- a/src/neptune/internal/operation_processors/async_operation_processor.py +++ b/src/neptune/internal/operation_processors/async_operation_processor.py @@ -51,7 +51,7 @@ from neptune.internal.operation_processors.operations_errors_processor import OperationsErrorsProcessor from neptune.internal.operation_processors.utils import common_metadata from neptune.internal.threading.daemon import Daemon -from neptune.internal.utils.disk_full import ensure_disk_not_full +from neptune.internal.utils.disk_utilization import ensure_disk_not_overutilize from neptune.internal.utils.logger import logger from neptune.internal.utils.monotonic_inc_batch_size import MonotonicIncBatchSize @@ -121,7 +121,7 @@ def _init_data_path(container_id: "UniqueId", container_type: "ContainerType") - path_suffix = f"exec-{now.timestamp()}-{now.strftime('%Y-%m-%d_%H.%M.%S.%f')}-{os.getpid()}" return get_container_dir(ASYNC_DIRECTORY, container_id, container_type, path_suffix) - @ensure_disk_not_full + @ensure_disk_not_overutilize def enqueue_operation(self, op: Operation, *, wait: bool) -> None: self._last_version = self._queue.put(op) diff --git a/src/neptune/internal/operation_processors/offline_operation_processor.py b/src/neptune/internal/operation_processors/offline_operation_processor.py index 6d3814ccb..3a31a567d 100644 --- a/src/neptune/internal/operation_processors/offline_operation_processor.py +++ b/src/neptune/internal/operation_processors/offline_operation_processor.py @@ -33,7 +33,7 @@ get_container_dir, ) from neptune.internal.operation_processors.utils import common_metadata -from neptune.internal.utils.disk_full import ensure_disk_not_full +from neptune.internal.utils.disk_utilization import ensure_disk_not_overutilize if TYPE_CHECKING: import threading @@ -60,7 +60,7 @@ def __init__(self, container_id: "UniqueId", container_type: "ContainerType", lo def _init_data_path(container_id: "UniqueId", container_type: "ContainerType") -> "Path": return get_container_dir(OFFLINE_DIRECTORY, container_id, container_type) - @ensure_disk_not_full + @ensure_disk_not_overutilize def enqueue_operation(self, op: Operation, *, wait: bool) -> None: self._queue.put(op) diff --git a/src/neptune/internal/operation_processors/sync_operation_processor.py b/src/neptune/internal/operation_processors/sync_operation_processor.py index fb71a4e23..885fa1251 100644 --- a/src/neptune/internal/operation_processors/sync_operation_processor.py +++ b/src/neptune/internal/operation_processors/sync_operation_processor.py @@ -30,7 +30,7 @@ get_container_dir, ) from neptune.internal.operation_processors.utils import common_metadata -from neptune.internal.utils.disk_full import ensure_disk_not_full +from neptune.internal.utils.disk_utilization import ensure_disk_not_overutilize if TYPE_CHECKING: from pathlib import Path @@ -60,7 +60,7 @@ def _init_data_path(container_id: "UniqueId", container_type: "ContainerType") - process_path = f"exec-{now.timestamp()}-{now.strftime('%Y-%m-%d_%H.%M.%S.%f')}-{os.getpid()}" return get_container_dir(SYNC_DIRECTORY, container_id, container_type, process_path) - @ensure_disk_not_full + @ensure_disk_not_overutilize def enqueue_operation(self, op: "Operation", *, wait: bool) -> None: _, errors = self._backend.execute_operations( container_id=self._container_id, diff --git a/src/neptune/internal/utils/disk_full.py b/src/neptune/internal/utils/disk_full.py deleted file mode 100644 index 64265981d..000000000 --- a/src/neptune/internal/utils/disk_full.py +++ /dev/null @@ -1,88 +0,0 @@ -# -# 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__ = ["ensure_disk_not_full"] - - -import os -from functools import wraps -from typing import ( - Any, - Callable, - Dict, - Optional, - Tuple, -) - -import psutil -from psutil import Error - -from neptune.common.warnings import ( - NeptuneWarning, - warn_once, -) -from neptune.constants import NEPTUNE_DATA_DIRECTORY -from neptune.envs import ( - NEPTUNE_MAX_DISK_UTILIZATION, - NEPTUNE_NON_RAISING_ON_DISK_ISSUE, -) - - -def get_neptune_data_directory() -> str: - return os.getenv("NEPTUNE_DATA_DIRECTORY", NEPTUNE_DATA_DIRECTORY) - - -def get_disk_utilization_percent(path: Optional[str] = None) -> float: - try: - if path is None: - path = get_neptune_data_directory() - - return float(psutil.disk_usage(path).percent) - except (ValueError, UnicodeEncodeError): - return 0 - - -def get_max_percentage_from_env() -> Optional[float]: - value = os.getenv(NEPTUNE_MAX_DISK_UTILIZATION) - if value is not None: - return float(value) - return None - - -def ensure_disk_not_full(func: Callable[..., None]) -> Callable[..., None]: - non_raising_on_disk_issue = NEPTUNE_NON_RAISING_ON_DISK_ISSUE in os.environ - max_disk_utilization = get_max_percentage_from_env() - - @wraps(func) - def wrapper(*args: Tuple, **kwargs: Dict[str, Any]) -> None: - if non_raising_on_disk_issue: - try: - if max_disk_utilization: - current_utilization = get_disk_utilization_percent() - 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.", - exception=NeptuneWarning, - ) - return - - func(*args, **kwargs) - except (OSError, Error): - warn_once("Encountered disk issue and Neptune will not be saving your data.", exception=NeptuneWarning) - else: - return func(*args, **kwargs) - - return wrapper diff --git a/src/neptune/internal/utils/disk_utilization.py b/src/neptune/internal/utils/disk_utilization.py new file mode 100644 index 000000000..ec3b74e08 --- /dev/null +++ b/src/neptune/internal/utils/disk_utilization.py @@ -0,0 +1,102 @@ +# +# 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__ = ["ensure_disk_not_overutilize"] + + +import os +from functools import wraps +from typing import ( + Any, + Callable, + Dict, + Optional, + Tuple, +) + +import psutil +from psutil import Error + +from neptune.common.warnings import ( + NeptuneWarning, + warn_once, +) +from neptune.constants import NEPTUNE_DATA_DIRECTORY +from neptune.envs import NEPTUNE_LIMIT_DISK_UTILIZATION + + +def get_neptune_data_directory() -> Optional[str]: + return os.getenv("NEPTUNE_DATA_DIRECTORY", NEPTUNE_DATA_DIRECTORY) + + +def get_disk_utilization_percent(path: Optional[str] = None) -> Optional[float]: + try: + if path is None: + path = get_neptune_data_directory() + if path is None: + return None + + return float(psutil.disk_usage(path).percent) + except (ValueError, TypeError, OSError, Error): + return None + + +def get_max_disk_utilization_from_env() -> Optional[float]: + env_limit_disk_utilization = os.getenv(NEPTUNE_LIMIT_DISK_UTILIZATION, "false") + + if env_limit_disk_utilization.lower() in ("false", "f", 0): + return None + + try: + limit_disk_utilization = float(env_limit_disk_utilization) + assert 0 < limit_disk_utilization <= 100 + + return limit_disk_utilization + except (ValueError, TypeError, AssertionError): + warn_once( + f"Provided invalid value of '{NEPTUNE_LIMIT_DISK_UTILIZATION}': '{env_limit_disk_utilization}'. " + "Check of disk utilization will not be applied.", + exception=NeptuneWarning, + ) + return None + + +def ensure_disk_not_overutilize(func: Callable[..., None]) -> Callable[..., None]: + max_disk_utilization = get_max_disk_utilization_from_env() + + @wraps(func) + def wrapper(*args: Tuple, **kwargs: Dict[str, Any]) -> None: + if max_disk_utilization: + current_utilization = get_disk_utilization_percent() + if current_utilization is None: + warn_once( + "Encountered disk issue during utilization check. Neptune will not save your data.", + exception=NeptuneWarning, + ) + return + + if current_utilization >= max_disk_utilization: + warn_once( + f"Max disk utilization {max_disk_utilization}% exceeded with {current_utilization}. " + + "Neptune will not be save your data.", + exception=NeptuneWarning, + ) + return + + func(*args, **kwargs) + else: + return func(*args, **kwargs) + + return wrapper diff --git a/tests/unit/neptune/new/internal/utils/test_disk_utilization.py b/tests/unit/neptune/new/internal/utils/test_disk_utilization.py new file mode 100644 index 000000000..68caa357d --- /dev/null +++ b/tests/unit/neptune/new/internal/utils/test_disk_utilization.py @@ -0,0 +1,100 @@ +# +# 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 +import warnings +from io import UnsupportedOperation + +from mock import ( + MagicMock, + patch, +) +from psutil import ( + AccessDenied, + Error, +) + +from neptune.envs import NEPTUNE_LIMIT_DISK_UTILIZATION +from neptune.internal.utils.disk_utilization import ensure_disk_not_overutilize + + +class TestDiskUtilization(unittest.TestCase): + def test_handle_invalid_env_values(self): + for value in ["True", "101", "-1"]: + with patch.dict(os.environ, {NEPTUNE_LIMIT_DISK_UTILIZATION: value}, clear=True): + mocked_func = MagicMock() + with warnings.catch_warnings(record=True) as warns: + wrapped_func = ensure_disk_not_overutilize(mocked_func) + wrapped_func() + + assert len(warns) == 1 + assert f"invalid value of '{NEPTUNE_LIMIT_DISK_UTILIZATION}': '{value}" in str(warns[-1].message) + mocked_func.assert_called_once() + + # 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_LIMIT_DISK_UTILIZATION: "60"}) + @patch("psutil.disk_usage") + def test_suppressing_of_env_errors(self, disk_usage_mock): + env_errors = [ + TypeError(), + OSError(), + IOError(), + EnvironmentError(), + UnsupportedOperation(), + Error(), + AccessDenied(), + ] + for error in env_errors: + mocked_func = MagicMock() + wrapped_func = ensure_disk_not_overutilize(mocked_func) + disk_usage_mock.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 = [OverflowError(), AttributeError()] + for error in non_env_errors: + mocked_func = MagicMock() + wrapped_func = ensure_disk_not_overutilize(mocked_func) + disk_usage_mock.side_effect = error + + with self.assertRaises(BaseException): + wrapped_func() + mocked_func.assert_not_called() + + @patch.dict(os.environ, {NEPTUNE_LIMIT_DISK_UTILIZATION: "100"}) + @patch("psutil.disk_usage") + def test_not_called_with_usage_100_percent(self, disk_usage_mock): + disk_usage_mock.return_value.percent = 100 + mocked_func = MagicMock() + wrapped_func = ensure_disk_not_overutilize(mocked_func) + + wrapped_func() + + mocked_func.assert_not_called() + + @patch.dict(os.environ, {NEPTUNE_LIMIT_DISK_UTILIZATION: "100"}) + @patch("psutil.disk_usage") + def test_called_when_usage_less_than_limit(self, disk_usage_mock): + disk_usage_mock.return_value.percent = 99 + mocked_func = MagicMock() + wrapped_func = ensure_disk_not_overutilize(mocked_func) + + wrapped_func() + + mocked_func.assert_called_once() diff --git a/tests/unit/neptune/new/internal/utils/test_full_disk.py b/tests/unit/neptune/new/internal/utils/test_full_disk.py deleted file mode 100644 index 2fac3c03e..000000000 --- a/tests/unit/neptune/new/internal/utils/test_full_disk.py +++ /dev/null @@ -1,87 +0,0 @@ -# -# 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() From 469413ae1ef03e3c222721f6156b352ac76fcc21 Mon Sep 17 00:00:00 2001 From: Artsiom Tserashkovich Date: Tue, 7 Nov 2023 15:58:54 +0100 Subject: [PATCH 2/6] Update src/neptune/internal/utils/disk_utilization.py Co-authored-by: Sabine --- src/neptune/internal/utils/disk_utilization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/neptune/internal/utils/disk_utilization.py b/src/neptune/internal/utils/disk_utilization.py index ec3b74e08..c29abfb20 100644 --- a/src/neptune/internal/utils/disk_utilization.py +++ b/src/neptune/internal/utils/disk_utilization.py @@ -89,8 +89,8 @@ def wrapper(*args: Tuple, **kwargs: Dict[str, Any]) -> None: if current_utilization >= max_disk_utilization: warn_once( - f"Max disk utilization {max_disk_utilization}% exceeded with {current_utilization}. " - + "Neptune will not be save your data.", + f"Disk usage is at {current_utilization}%, which exceeds the maximum allowed utilization " + + "of {max_disk_utilization}%. Neptune will not save your data.", exception=NeptuneWarning, ) return From dd9622c3a056c605640f095e56cb705e72dd1788 Mon Sep 17 00:00:00 2001 From: Artsiom Tserashkovich Date: Tue, 7 Nov 2023 17:11:04 +0100 Subject: [PATCH 3/6] Return 2 env vars for control disk usage --- src/neptune/envs.py | 7 +++-- .../internal/utils/disk_utilization.py | 18 ++++++----- .../internal/utils/test_disk_utilization.py | 31 +++++++++++++++---- 3 files changed, 40 insertions(+), 16 deletions(-) diff --git a/src/neptune/envs.py b/src/neptune/envs.py index c262b8413..df9df83ff 100644 --- a/src/neptune/envs.py +++ b/src/neptune/envs.py @@ -31,7 +31,8 @@ "NEPTUNE_FETCH_TABLE_STEP_SIZE", "NEPTUNE_SYNC_AFTER_STOP_TIMEOUT", "NEPTUNE_REQUEST_TIMEOUT", - "NEPTUNE_LIMIT_DISK_UTILIZATION", + "NEPTUNE_MAX_DISK_UTILIZATION", + "NEPTUNE_NON_RAISING_ON_DISK_ISSUE", "NEPTUNE_ENABLE_DEFAULT_ASYNC_LAG_CALLBACK", "NEPTUNE_ENABLE_DEFAULT_ASYNC_NO_PROGRESS_CALLBACK", "NEPTUNE_DISABLE_PARENT_DIR_DELETION", @@ -77,7 +78,9 @@ NEPTUNE_ENABLE_DEFAULT_ASYNC_NO_PROGRESS_CALLBACK = "NEPTUNE_ENABLE_DEFAULT_ASYNC_NO_PROGRESS_CALLBACK" -NEPTUNE_LIMIT_DISK_UTILIZATION = "NEPTUNE_LIMIT_DISK_UTILIZATION" +NEPTUNE_MAX_DISK_UTILIZATION = "NEPTUNE_MAX_DISK_UTILIZATION" + +NEPTUNE_NON_RAISING_ON_DISK_ISSUE = "NEPTUNE_NON_RAISING_ON_DISK_ISSUE" NEPTUNE_DISABLE_PARENT_DIR_DELETION = "NEPTUNE_DISABLE_PARENT_DIR_DELETION" diff --git a/src/neptune/internal/utils/disk_utilization.py b/src/neptune/internal/utils/disk_utilization.py index c29abfb20..277a01c9c 100644 --- a/src/neptune/internal/utils/disk_utilization.py +++ b/src/neptune/internal/utils/disk_utilization.py @@ -34,10 +34,13 @@ warn_once, ) from neptune.constants import NEPTUNE_DATA_DIRECTORY -from neptune.envs import NEPTUNE_LIMIT_DISK_UTILIZATION +from neptune.envs import ( + NEPTUNE_MAX_DISK_UTILIZATION, + NEPTUNE_NON_RAISING_ON_DISK_ISSUE, +) -def get_neptune_data_directory() -> Optional[str]: +def get_neptune_data_directory() -> str: return os.getenv("NEPTUNE_DATA_DIRECTORY", NEPTUNE_DATA_DIRECTORY) @@ -45,8 +48,6 @@ def get_disk_utilization_percent(path: Optional[str] = None) -> Optional[float]: try: if path is None: path = get_neptune_data_directory() - if path is None: - return None return float(psutil.disk_usage(path).percent) except (ValueError, TypeError, OSError, Error): @@ -54,9 +55,9 @@ def get_disk_utilization_percent(path: Optional[str] = None) -> Optional[float]: def get_max_disk_utilization_from_env() -> Optional[float]: - env_limit_disk_utilization = os.getenv(NEPTUNE_LIMIT_DISK_UTILIZATION, "false") + env_limit_disk_utilization = os.getenv(NEPTUNE_MAX_DISK_UTILIZATION) - if env_limit_disk_utilization.lower() in ("false", "f", 0): + if env_limit_disk_utilization is None: return None try: @@ -66,7 +67,7 @@ def get_max_disk_utilization_from_env() -> Optional[float]: return limit_disk_utilization except (ValueError, TypeError, AssertionError): warn_once( - f"Provided invalid value of '{NEPTUNE_LIMIT_DISK_UTILIZATION}': '{env_limit_disk_utilization}'. " + f"Provided invalid value of '{NEPTUNE_MAX_DISK_UTILIZATION}': '{env_limit_disk_utilization}'. " "Check of disk utilization will not be applied.", exception=NeptuneWarning, ) @@ -74,11 +75,12 @@ def get_max_disk_utilization_from_env() -> Optional[float]: def ensure_disk_not_overutilize(func: Callable[..., None]) -> Callable[..., None]: + non_raising_on_disk_issue = os.getenv(NEPTUNE_NON_RAISING_ON_DISK_ISSUE, "false").lower() in ("true", "t", "1") max_disk_utilization = get_max_disk_utilization_from_env() @wraps(func) def wrapper(*args: Tuple, **kwargs: Dict[str, Any]) -> None: - if max_disk_utilization: + if non_raising_on_disk_issue and max_disk_utilization: current_utilization = get_disk_utilization_percent() if current_utilization is None: warn_once( diff --git a/tests/unit/neptune/new/internal/utils/test_disk_utilization.py b/tests/unit/neptune/new/internal/utils/test_disk_utilization.py index 68caa357d..29db36550 100644 --- a/tests/unit/neptune/new/internal/utils/test_disk_utilization.py +++ b/tests/unit/neptune/new/internal/utils/test_disk_utilization.py @@ -27,27 +27,32 @@ Error, ) -from neptune.envs import NEPTUNE_LIMIT_DISK_UTILIZATION +from neptune.envs import ( + NEPTUNE_MAX_DISK_UTILIZATION, + NEPTUNE_NON_RAISING_ON_DISK_ISSUE, +) from neptune.internal.utils.disk_utilization import ensure_disk_not_overutilize class TestDiskUtilization(unittest.TestCase): + @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"}) def test_handle_invalid_env_values(self): for value in ["True", "101", "-1"]: - with patch.dict(os.environ, {NEPTUNE_LIMIT_DISK_UTILIZATION: value}, clear=True): + with patch.dict(os.environ, {NEPTUNE_MAX_DISK_UTILIZATION: value}, clear=True): mocked_func = MagicMock() with warnings.catch_warnings(record=True) as warns: wrapped_func = ensure_disk_not_overutilize(mocked_func) wrapped_func() assert len(warns) == 1 - assert f"invalid value of '{NEPTUNE_LIMIT_DISK_UTILIZATION}': '{value}" in str(warns[-1].message) + assert f"invalid value of '{NEPTUNE_MAX_DISK_UTILIZATION}': '{value}" in str(warns[-1].message) mocked_func.assert_called_once() # 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_LIMIT_DISK_UTILIZATION: "60"}) + @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"}) + @patch.dict(os.environ, {NEPTUNE_MAX_DISK_UTILIZATION: "60"}) @patch("psutil.disk_usage") def test_suppressing_of_env_errors(self, disk_usage_mock): env_errors = [ @@ -77,7 +82,8 @@ def test_suppressing_of_env_errors(self, disk_usage_mock): wrapped_func() mocked_func.assert_not_called() - @patch.dict(os.environ, {NEPTUNE_LIMIT_DISK_UTILIZATION: "100"}) + @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"}) + @patch.dict(os.environ, {NEPTUNE_MAX_DISK_UTILIZATION: "100"}) @patch("psutil.disk_usage") def test_not_called_with_usage_100_percent(self, disk_usage_mock): disk_usage_mock.return_value.percent = 100 @@ -88,7 +94,8 @@ def test_not_called_with_usage_100_percent(self, disk_usage_mock): mocked_func.assert_not_called() - @patch.dict(os.environ, {NEPTUNE_LIMIT_DISK_UTILIZATION: "100"}) + @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"}) + @patch.dict(os.environ, {NEPTUNE_MAX_DISK_UTILIZATION: "100"}) @patch("psutil.disk_usage") def test_called_when_usage_less_than_limit(self, disk_usage_mock): disk_usage_mock.return_value.percent = 99 @@ -98,3 +105,15 @@ def test_called_when_usage_less_than_limit(self, disk_usage_mock): wrapped_func() mocked_func.assert_called_once() + + @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "False"}) + @patch.dict(os.environ, {NEPTUNE_MAX_DISK_UTILIZATION: "60"}) + @patch("psutil.disk_usage") + def test_not_called_when_(self, disk_usage_mock): + disk_usage_mock.return_value.percent = 99 + mocked_func = MagicMock() + wrapped_func = ensure_disk_not_overutilize(mocked_func) + + wrapped_func() + + mocked_func.assert_called_once() From 6c0a874fae072f54d6202939be156d3a8cdf04b8 Mon Sep 17 00:00:00 2001 From: Artsiom Tserashkovich Date: Wed, 8 Nov 2023 10:48:43 +0100 Subject: [PATCH 4/6] Rework condition to follow previous flow --- .../internal/utils/disk_utilization.py | 43 +++++++++++-------- .../internal/utils/test_disk_utilization.py | 37 +++++++++++----- 2 files changed, 50 insertions(+), 30 deletions(-) diff --git a/src/neptune/internal/utils/disk_utilization.py b/src/neptune/internal/utils/disk_utilization.py index 277a01c9c..cf38752c4 100644 --- a/src/neptune/internal/utils/disk_utilization.py +++ b/src/neptune/internal/utils/disk_utilization.py @@ -50,7 +50,7 @@ def get_disk_utilization_percent(path: Optional[str] = None) -> Optional[float]: path = get_neptune_data_directory() return float(psutil.disk_usage(path).percent) - except (ValueError, TypeError, OSError, Error): + except (ValueError, TypeError, Error): return None @@ -80,24 +80,29 @@ def ensure_disk_not_overutilize(func: Callable[..., None]) -> Callable[..., None @wraps(func) def wrapper(*args: Tuple, **kwargs: Dict[str, Any]) -> None: - if non_raising_on_disk_issue and max_disk_utilization: - current_utilization = get_disk_utilization_percent() - if current_utilization is None: - warn_once( - "Encountered disk issue during utilization check. Neptune will not save your data.", - exception=NeptuneWarning, - ) - return - - if current_utilization >= max_disk_utilization: - warn_once( - f"Disk usage is at {current_utilization}%, which exceeds the maximum allowed utilization " - + "of {max_disk_utilization}%. Neptune will not save your data.", - exception=NeptuneWarning, - ) - return - - func(*args, **kwargs) + if non_raising_on_disk_issue: + try: + if max_disk_utilization: + current_utilization = get_disk_utilization_percent() + if current_utilization is None: + warn_once( + "Encountered disk issue during utilization check. Neptune will not save your data.", + exception=NeptuneWarning, + ) + return + + if current_utilization >= max_disk_utilization: + warn_once( + f"Disk usage is at {current_utilization}%, which exceeds the maximum allowed utilization " + + "of {max_disk_utilization}%. Neptune will not save your data.", + exception=NeptuneWarning, + ) + return + + func(*args, **kwargs) + except (OSError, Error): + warn_once("Encountered disk issue. Neptune will not save your data.", exception=NeptuneWarning) + else: return func(*args, **kwargs) diff --git a/tests/unit/neptune/new/internal/utils/test_disk_utilization.py b/tests/unit/neptune/new/internal/utils/test_disk_utilization.py index 29db36550..40a690654 100644 --- a/tests/unit/neptune/new/internal/utils/test_disk_utilization.py +++ b/tests/unit/neptune/new/internal/utils/test_disk_utilization.py @@ -52,11 +52,8 @@ def test_handle_invalid_env_values(self): # 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.dict(os.environ, {NEPTUNE_MAX_DISK_UTILIZATION: "60"}) - @patch("psutil.disk_usage") - def test_suppressing_of_env_errors(self, disk_usage_mock): - env_errors = [ - TypeError(), + def test_suppressing_of_func_errors(self): + disk_errors = [ OSError(), IOError(), EnvironmentError(), @@ -64,22 +61,40 @@ def test_suppressing_of_env_errors(self, disk_usage_mock): Error(), AccessDenied(), ] - for error in env_errors: + for error in disk_errors: mocked_func = MagicMock() wrapped_func = ensure_disk_not_overutilize(mocked_func) - disk_usage_mock.side_effect = error + mocked_func.side_effect = error wrapped_func() # asserting is not required as expecting that any error will be caught - mocked_func.assert_not_called() + mocked_func.assert_called_once() - non_env_errors = [OverflowError(), AttributeError()] - for error in non_env_errors: + non_disk_errors = [OverflowError(), AttributeError()] + for error in non_disk_errors: mocked_func = MagicMock() wrapped_func = ensure_disk_not_overutilize(mocked_func) - disk_usage_mock.side_effect = error + mocked_func.side_effect = error with self.assertRaises(BaseException): wrapped_func() + mocked_func.assert_called_once() + + @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"}) + @patch.dict(os.environ, {NEPTUNE_MAX_DISK_UTILIZATION: "60"}) + @patch("psutil.disk_usage") + def test_suppressing_of_checking_utilization_errors(self, disk_usage_mock): + checking_errors = [ + TypeError(), + UnsupportedOperation(), + Error(), + AccessDenied(), + ] + for error in checking_errors: + mocked_func = MagicMock() + wrapped_func = ensure_disk_not_overutilize(mocked_func) + disk_usage_mock.side_effect = error + + wrapped_func() # asserting is not required as expecting that any error will be caught mocked_func.assert_not_called() @patch.dict(os.environ, {NEPTUNE_NON_RAISING_ON_DISK_ISSUE: "True"}) From 2506f8c02b159e469cb4b790d103d4da1716bf2a Mon Sep 17 00:00:00 2001 From: Artsiom Tserashkovich Date: Wed, 8 Nov 2023 15:22:31 +0100 Subject: [PATCH 5/6] Use ValueErro instead of assert --- src/neptune/internal/utils/disk_utilization.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/neptune/internal/utils/disk_utilization.py b/src/neptune/internal/utils/disk_utilization.py index cf38752c4..a0d138126 100644 --- a/src/neptune/internal/utils/disk_utilization.py +++ b/src/neptune/internal/utils/disk_utilization.py @@ -62,10 +62,11 @@ def get_max_disk_utilization_from_env() -> Optional[float]: try: limit_disk_utilization = float(env_limit_disk_utilization) - assert 0 < limit_disk_utilization <= 100 + if limit_disk_utilization <= 0 or limit_disk_utilization > 100: + raise ValueError return limit_disk_utilization - except (ValueError, TypeError, AssertionError): + except (ValueError, TypeError): warn_once( f"Provided invalid value of '{NEPTUNE_MAX_DISK_UTILIZATION}': '{env_limit_disk_utilization}'. " "Check of disk utilization will not be applied.", From abd216660c61caf64e46f3b7e10015d1088ffe0b Mon Sep 17 00:00:00 2001 From: Artsiom Tserashkovich Date: Wed, 8 Nov 2023 15:47:28 +0100 Subject: [PATCH 6/6] update CHANGELOG.md --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c9f673cb..7b0a8e41e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ - Minor tweaks to `neptune.cli` and cleaning leftovers after async Experiments ([#1529](https://github.com/neptune-ai/neptune-client/pull/1529)) - Pin `simplejson` required version to below `3.19` ([#1535](https://github.com/neptune-ai/neptune-client/pull/1535)) - Added `experimental` mode that supports partitioned operations queue ([#1524](https://github.com/neptune-ai/neptune-client/pull/1524)) +- Rework disk utilization check ([#1549](https://github.com/neptune-ai/neptune-client/pull/1549)) ## neptune 1.8.2