Repository: caikit/caikit (python)

Issue:
Unsafe rmtree in ModuleSaver on any error can cause data loss :bomb:
## Describe the bug

ModuleSaver wants to rmtree to remove any files it created (or might have touched) if any error occurs.  This is assuming that the given model_path was either newly created or no longer usable due to some saver error (harsh but maybe appropriate in specific use cases).

But! model_path may contain important files that should not be deleted.  In the worst case, if the error was caused by an inappropriate model_path (e.g. "" or "/") then the _exit will just wipeout files.  This would not only clobber some model directory, or files created/touched, but potentially anything on your disk.

Even if the error was on a good model_path with existing files, reverting to prior state would be more appropriate than :bomb:

## Platform

## Sample Code

don't try it

## Expected behavior

shutil.rmtree() is dangerous and should never be used except with tightly controlled parameters.

## Observed behavior

My local caikit-nlp repo was wiped out, when testing without setting a model_path.  I think this is what caused it.  Just a bootstrap/save script that lacked arg checks (which is apparently how folks are supposed to get model config files).

## Additional context

The general behavior of rmtree is the bigger concern (please don't), but the use case here where ModuleSaver seems to be designed with exist_ok=True and yet expects to wipeout that path when errors occur is questionable.  This really needs to be able to revert any changed file on error.

Pull request:
Add exist_ok option to ModuleSaver and only rmtree when exist_ok=False
The prior implementation was dangerous because it could remove existing files including existing models or any accidental path passed in. This is a bit safer.

Fixes #525

<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility


Patch:
diff --git a/caikit/core/modules/saver.py b/caikit/core/modules/saver.py
index 7007463b..a6d70abe 100644
--- a/caikit/core/modules/saver.py
+++ b/caikit/core/modules/saver.py
@@ -41,7 +41,8 @@ error = error_handler.get(log)
 
 class ModuleSaver:
     """A module saver that provides common functionality used for saving modules and also a context
-    manager that cleans up gracefully in case an error is encountered during the save process.
+    manager that cleans up in case an error is encountered during the save process for a model_path
+    that did not already exist.
     """
 
     SAVED_KEY_NAME = "saved"
@@ -51,7 +52,7 @@ class ModuleSaver:
     MODULE_ID_KEY_NAME = "module_id"
     MODULE_CLASS_KEY_NAME = "module_class"
 
-    def __init__(self, module: ModuleBase, model_path):
+    def __init__(self, module: ModuleBase, model_path, exist_ok=True):
         """Construct a new module saver.
 
         Args:
@@ -60,8 +61,10 @@ class ModuleSaver:
             model_path (str): The absolute path to the directory where the model
                 will be saved.  If this directory does not exist, it will be
                 created.
+            exist_ok (bool): Allow to overwrite existing model_path files.
         """
         self.model_path = os.path.normpath(model_path)
+        self.exist_ok = exist_ok
 
         # Get possibly nested caikit library path
         module_path = module.__module__
@@ -313,20 +316,31 @@ class ModuleSaver:
     def __enter__(self):
         """Enter the module saver context.  This creates the `model_path` directory.  If this
         context successfully exits, then the model configuration and all files it contains will
-        be written and saved to disk inside the `model_path` directory.  If any uncaught exceptions
-        are thrown inside this context, then `model_path` will be removed.
+        be written and saved to disk inside the `model_path` directory.
+
+        If `exist_ok` is False, an exception will be raised before touching existing `model_path`
+        files.
+
+        If any uncaught exceptions are thrown inside this context, and `exist_ok` is False,
+        then this new `model_path` will be removed. If `exist_ok` is True, the files will be kept
+        and may include incomplete updates.
         """
-        os.makedirs(self.model_path, exist_ok=True)
+        os.makedirs(self.model_path, exist_ok=self.exist_ok)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Exit the module saver context. If this context successfully exits, then the model
         configuration and all files it contains will be written and saved to disk inside the
-        `model_path` directory.  If any uncaught exceptions are thrown inside this context, then
-        `model_path` will be removed.
+        `model_path` directory.
+
+        If any uncaught exceptions are thrown inside this context, and `exist_ok` is False,
+        then this new `model_path` will be removed. If `exist_ok` is True, the files will be kept
+        and may include incomplete updates.
         """
         if exc_type is not None:
-            shutil.rmtree(self.model_path, ignore_errors=True)
+            if not self.exist_ok:
+                # Presume it is okay to rmtree
+                shutil.rmtree(self.model_path, ignore_errors=True)
             return
 
         ModuleConfig(self.config).save(self.model_path)


Test patch:
diff --git a/tests/core/modules/test_module.py b/tests/core/modules/test_module.py
index eaab139f..0ecaba7c 100644
--- a/tests/core/modules/test_module.py
+++ b/tests/core/modules/test_module.py
@@ -156,6 +156,61 @@ def test_save_module(model_path):
             assert module_saver.config.get("module_paths") == {"dummy": "./dummy"}
 
 
+def test_save_please_dont_destroy(model_path):
+    dummy_model = caikit.core.load(model_path)
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        # exist_ok=False raises FileExistsError
+        with pytest.raises(FileExistsError):
+            with ModuleSaver(
+                dummy_model,
+                model_path=tempdir,
+                exist_ok=False,
+            ):
+                pass
+
+        # Existing dir should definitely not be removed
+        assert os.path.exists(tempdir)
+
+        # exist_ok=True does not raise error
+        with ModuleSaver(
+            dummy_model,
+            model_path=tempdir,
+            exist_ok=True,
+        ):
+            pass
+        assert os.path.exists(tempdir)
+
+        # exist_ok=True and exception thrown -> please don't destroy!
+        with pytest.raises(ValueError):
+            with ModuleSaver(
+                dummy_model,
+                model_path=tempdir,
+                exist_ok=True,
+            ):
+                raise ValueError  # a test exception
+        assert os.path.exists(tempdir)
+
+
+def test_save_okay_to_destroy(model_path):
+    dummy_model = caikit.core.load(model_path)
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        test_subdir = os.path.join(tempdir, "subdir")
+
+        # exist_ok=False and exception thrown -> okay to destroy
+        with pytest.raises(ValueError):
+            with ModuleSaver(
+                dummy_model,
+                model_path=test_subdir,
+                exist_ok=False,
+            ):
+                assert os.path.exists(test_subdir)
+                raise ValueError  # a test exception
+
+        assert not os.path.exists(test_subdir)
+
+
 def test_save_module_kwargs_get_piped_through(model_path):
     """Testing additional keywords passed to save_module are processed by the module"""
     dummy_model = caikit.core.load(model_path)