By default avoiding generating files in temp directory (#1609)

* initial commit

* all pathways validated, aiming to get tests green

* all pathways validated, aiming to get tests green

* Apply suggestions from code review

* Update src/sparseml/yolov8/trainers.py

* Update src/sparseml/onnx/optim/quantization/calibration.py

* torch and onnx tests green

* fix quality

* add print to conftest to see which file is the culprit

* improved diagnosis

* make sure no files are being created

* update conftest

* make conftest faster

* quality

* ready for review

* revert to keep a benign script

* Update torch_to_onnx_exporter.py

* Revert "Update torch_to_onnx_exporter.py"

This reverts commit 6dbf554.

---------

Co-authored-by: bogunowicz@arrival.com <bogunowicz@arrival.com>

src/sparseml/onnx/optim/quantization/calibration.py

@@ -17,8 +17,8 @@
 """
 
 
+import logging
 import os
-import tempfile
 from typing import Dict, Generator, Iterable, List, Tuple, Union
 
 import numpy as np
@@ -28,6 +28,8 @@
 from sparsezoo.utils import save_onnx, validate_onnx
 
 
+_LOGGER = logging.getLogger(__name__)
+
 __all__ = ["CalibrationSession"]
 
 
@@ -65,11 +67,11 @@ def __init__(
         self._model_augmented = self.generate_augmented_model()
 
         if self._augmented_model_path is None:
-            self._augmented_model_tmp_file = tempfile.NamedTemporaryFile(
-                suffix=".onnx", delete=True
+            self._augmented_model_path = os.path.join(
+                os.getcwd(), "model_augmented.onnx"
             )
-            self._augmented_model_path = self._augmented_model_tmp_file.name
         save_onnx(self._model_augmented, self._augmented_model_path)
+        _LOGGER.debug(f"Created an augmented model at: {self._augmented_model_path}")
 
         self._sessions = {}  # batch_size -> session
         self._quantization_thresholds = {}  # Dict[node.name, Tuple(min_val, max_val)]
@@ -103,13 +105,15 @@ def _optimize_model(self) -> Union[str, None]:
                 # no optimization performed, skip the rest of this block
                 raise Exception()
             validate_onnx(model_optimized)  # should raise exception if broken
-            optimized_model_path = tempfile.NamedTemporaryFile(
-                suffix=".onnx", delete=False
-            )
-            save_onnx(model_optimized, optimized_model_path.name)
+            optimized_model_path = os.path.join(os.getcwd(), "model_optimized.onnx")
+            save_onnx(model_optimized, optimized_model_path)
             self._model = model_optimized
-            print("Optimization successful")
-            return optimized_model_path.name
+            _LOGGER.debug(
+                "Optimization successful. "
+                "Created an optimized model at: "
+                f"{optimized_model_path}"
+            )
+            return optimized_model_path
         except Exception as e:
             print(e)
             print(
@@ -352,3 +356,5 @@ def __del__(self):
         """
         if self._optimized_model_path is not None:
             os.remove(self._optimized_model_path)
+        if self._augmented_model_path is not None:
+            os.remove(self._augmented_model_path)

src/sparseml/pytorch/torch_to_onnx_exporter.py

@@ -14,7 +14,7 @@
 
 import collections
 import logging
-import tempfile
+import os
 import warnings
 from copy import deepcopy
 from typing import Any, Dict, Iterable, List
@@ -147,9 +147,10 @@ def post_validate(self, model: onnx.ModelProto) -> onnx.ModelProto:
             raise TypeError(f"Expected onnx.ModelProto, found {type(model)}")
         return model
 
-    def transform(self, module: torch.nn.Module) -> onnx.ModelProto:
-        tmp = tempfile.NamedTemporaryFile("w")
-        file_path = tmp.name
+    def transform(
+        self, module: torch.nn.Module, file_name: str = "model.onnx"
+    ) -> onnx.ModelProto:
+        file_path = os.path.join(os.getcwd(), file_name)
 
         _LOGGER.debug(f"Saving onnx model to {file_path}")
 

tests/integrations/image_classification/test_image_classification.py

@@ -107,7 +107,8 @@ def teardown(self):
         """
         Cleanup environment after test completion
         """
-        pass
+        if "train" in self.command_types:
+            self.save_dir.cleanup()
 
 
 class TestImageClassification(BaseIntegrationTester):

tests/integrations/transformers/test_transformers.py

@@ -142,7 +142,9 @@ def get_root_commands(self, raw_configs):
         return command_stubs_final
 
     def teardown(self):
-        pass  # not yet implemented
+        # clean up temporary directory
+        if "train" in self.command_types:
+            self.save_dir.cleanup()
 
 
 @flaky(max_runs=2, min_passes=1)

tests/sparseml/conftest.py

@@ -14,6 +14,8 @@
 
 import os
 import shutil
+import tempfile
+from typing import List
 
 import pytest
 
@@ -28,16 +30,52 @@
 os.environ["NM_TEST_LOG_DIR"] = "nm_temp_test_logs"
 
 
+def _get_files(directory: str) -> List[str]:
+    list_filepaths = []
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            list_filepaths.append(os.path.join(os.path.abspath(root), file))
+    return list_filepaths
+
+
 @pytest.fixture(scope="session", autouse=True)
 def check_for_created_files():
-    start_file_count = sum(len(files) for _, _, files in os.walk(r"."))
+    start_files_root = _get_files(directory=r".")
+    start_files_temp = _get_files(directory=tempfile.gettempdir())
     yield
     if wandb:
         wandb.finish()
     log_dir = os.environ.get("NM_TEST_LOG_DIR")
     if os.path.isdir(log_dir):
         shutil.rmtree(log_dir)
-    end_file_count = sum(len(files) for _, _, files in os.walk(r"."))
-    assert (
-        start_file_count >= end_file_count
-    ), f"{end_file_count - start_file_count} files created during pytest run"
+
+    end_files_root = _get_files(directory=r".")
+    end_files_temp = _get_files(directory=tempfile.gettempdir())
+
+    # assert no files created in root directory while running
+    # the pytest suite
+    assert len(start_files_root) >= len(end_files_root), (
+        f"{len(end_files_root) - len(start_files_root)} "
+        f"files created in current working "
+        f"directory during pytest run. "
+        f"Created files: {set(end_files_root) - set(start_files_root)}"
+    )
+    max_allowed_sized_temp_files_megabytes = 1
+    created_temp_files = set(end_files_temp) - set(start_files_temp)
+    size_of_temp_files_bytes = 0
+    for file_path in created_temp_files:
+        try:
+            size_of_temp_files_bytes += os.path.getsize(file_path)
+        # if file is deleted between the time we get the list of files
+        # and the time we get the size of the file, ignore it
+        except FileNotFoundError:
+            pass
+
+    size_of_temp_files_megabytes = size_of_temp_files_bytes / 1024 / 1024
+    # assert no more than 1 megabyte of temp files created in temp directory
+    # while running the pytest suite
+    assert max_allowed_sized_temp_files_megabytes >= size_of_temp_files_megabytes, (
+        f"{size_of_temp_files_megabytes} "
+        f"megabytes of temp files created in temp directory during pytest run. "
+        f"Created files: {set(end_files_temp) - set(start_files_temp)}"
+    )

tests/sparseml/exporters/transforms/test_onnx_transform.py

@@ -64,3 +64,4 @@ def test_onnx_transform_from_path(tmp_path):
     path = os.path.join(str(tmp_path), "model.onnx")
     save_onnx(model, path)
     assert transform(path)
+    os.remove(path)

tests/sparseml/optim/test_analyzer.py

@@ -23,11 +23,13 @@ def test_layer_descs():
         AnalyzedLayerDesc("layer1", "Linear"),
         AnalyzedLayerDesc("layer2", "Conv2d"),
     ]
-    save_path = os.path.join(tempfile.gettempdir(), "layer_descs.json")
+    tempdir = tempfile.gettempdir()
+    save_path = os.path.join(tempdir, "layer_descs.json")
 
     AnalyzedLayerDesc.save_descs(descs, save_path)
     loaded_descs = AnalyzedLayerDesc.load_descs(save_path)
 
     for desc, loaded_desc in zip(descs, loaded_descs):
         assert desc.name == loaded_desc.name
         assert desc.type_ == loaded_desc.type_
+    os.remove(save_path)

tests/sparseml/optim/test_sensitivity.py

@@ -143,7 +143,8 @@ def test_ks_loss_sensitivity_analysis():
         assert abs(test["sparse_integral"] - comp.sparse_integral) < 1e-5
         assert abs(test["sparse_comparison"] - comp.sparse_comparison()) < 1e-5
 
-    path = os.path.join(tempfile.gettempdir(), "ks-sens-analysis.json")
+    tempdir = tempfile.TemporaryDirectory()
+    path = os.path.join(tempdir.name, "ks-sens-analysis.json")
     analysis.save_json(path)
 
     json_analysis = analysis.load_json(path)
@@ -160,26 +161,24 @@ def test_ks_loss_sensitivity_analysis():
         assert abs(test["sparse_integral"] - comp.sparse_integral) < 1e-5
         assert abs(test["sparse_comparison"] - comp.sparse_comparison()) < 1e-5
 
-    path = os.path.join(
-        tempfile.gettempdir(), "ks-sens-analysis-integral-normalized.png"
-    )
+    path = os.path.join(tempdir.name, "ks-sens-analysis-integral-normalized.png")
     analysis.plot(path, plot_integral=True)
     assert os.path.exists(path)
 
-    path = os.path.join(
-        tempfile.gettempdir(), "ks-sens-analysis-integral-normalized.png"
-    )
+    path = os.path.join(tempdir.name, "ks-sens-analysis-integral-normalized.png")
     analysis.plot(path, plot_integral=True)
     assert os.path.exists(path)
 
-    path = os.path.join(tempfile.gettempdir(), "ks-sens-analysis-integral.png")
+    path = os.path.join(tempdir.name, "ks-sens-analysis-integral.png")
     analysis.plot(path, plot_integral=True, normalize=False)
     assert os.path.exists(path)
 
-    path = os.path.join(tempfile.gettempdir(), "ks-sens-analysis-avg-normalized.png")
+    path = os.path.join(tempdir.name, "ks-sens-analysis-avg-normalized.png")
     analysis.plot(path, plot_integral=False)
     assert os.path.exists(path)
 
-    path = os.path.join(tempfile.gettempdir(), "ks-sens-analysis-avg.png")
+    path = os.path.join(tempdir.name, "ks-sens-analysis-avg.png")
     analysis.plot(path, plot_integral=False, normalize=False)
     assert os.path.exists(path)
+
+    tempdir.cleanup()

tests/sparseml/pytorch/image_classification/test_export.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import math
+import shutil
 from pathlib import Path
 
 import onnx
@@ -96,3 +97,4 @@ def test_export_one_shot(resnet_checkpoint, recipe_path, temp_dir):
             node_sparsity = node.parameter_summary.block_structure["single"].sparsity
             assert math.isclose(node_sparsity, 0.9, abs_tol=0.01)
     assert found_prunable_node
+    shutil.rmtree(temp_dir)

tests/sparseml/pytorch/image_classification/utils/test_helpers.py

@@ -11,6 +11,8 @@
 # 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 shutil
+
 import pytest
 
 from sparseml.pytorch.image_classification.utils.helpers import save_zoo_directory
@@ -41,3 +43,5 @@ def test_save_zoo_directory(stub, tmp_path_factory):
     )
     new_zoo_model = Model(str(save_dir))
     assert new_zoo_model.validate(minimal_validation=True, validate_onnxruntime=False)
+    shutil.rmtree(path_to_training_outputs)
+    shutil.rmtree(save_dir)

tests/sparseml/pytorch/recipe_template/test_end_to_end.py

@@ -11,6 +11,7 @@
 # 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 shutil
 
 import pytest
 
@@ -45,3 +46,4 @@ def test_docstring_cli_examples(command, tmp_path):
     command.extend(["--file_name", str(tmp_path / "temp.md")])
     result = runner.invoke(main, command)
     assert result.exit_code == 0
+    shutil.rmtree(tmp_path)

tests/sparseml/pytorch/sparsification/test_modifier_thinning.py

@@ -237,3 +237,4 @@ def test_structured_pruning_one_shot_e2e(
         assert abs(applied_compression - sparsity) < 5e-2
     # validate forward pass
     assert module(torch.randn(2, 3, 224, 224)) is not None
+    os.remove(tmp_fp)

tests/sparseml/pytorch/test_torch_to_onnx_exporter.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import shutil
 from collections import Counter
 
 import numpy
@@ -53,6 +54,7 @@ def test_simple_models_against_module_exporter(tmp_path, model, sample_batch):
         str(tmp_path / "new_exporter" / "model.onnx"),
         sample_batch,
     )
+    shutil.rmtree(tmp_path)
 
 
 @pytest.mark.skipif(
@@ -105,6 +107,7 @@ def test_resnet50_exporters_are_equivalent(tmp_path, quantize: bool, convert_qat
         sample_batch,
         expect_op_types=["QuantizeLinear", "DequantizeLinear"] if quantize else None,
     )
+    shutil.rmtree(tmp_path)
 
 
 @pytest.mark.skipif(

tests/sparseml/pytorch/utils/test_exporter.py

@@ -36,8 +36,10 @@
 )
 def test_exporter_onnx():
     sample_batch = torch.randn(1, 8)
-    exporter = ModuleExporter(MLPNet(), tempfile.gettempdir())
+    tempdir = tempfile.TemporaryDirectory()
+    exporter = ModuleExporter(MLPNet(), tempdir.name)
     exporter.export_onnx(sample_batch)
+    tempdir.cleanup()
 
 
 @pytest.mark.skipif(
@@ -47,8 +49,11 @@ def test_exporter_onnx():
 @pytest.mark.parametrize("batch_size", [1, 64])
 def test_export_batches(batch_size):
     sample_batch = torch.randn(batch_size, 8)
-    exporter = ModuleExporter(MLPNet(), tempfile.gettempdir())
+    # create a directory in tempdir
+    tempdir = tempfile.TemporaryDirectory()
+    exporter = ModuleExporter(MLPNet(), tempdir.name)
     exporter.export_samples([sample_batch])
+    tempdir.cleanup()
 
 
 def test_fold_identity_initializers():

tests/sparseml/pytorch/utils/test_helpers.py

@@ -558,6 +558,7 @@ def test_tensor_export_npy(tensor, name):
 
     for s1, s2 in zip(exported.shape, tensor.shape):
         assert s1 == s2
+    os.remove(path)
 
 
 @pytest.mark.skipif(
@@ -579,6 +580,7 @@ def test_tensor_export_npz(tensor, name):
 
     for s1, s2 in zip(exported.shape, tensor.shape):
         assert s1 == s2
+    os.remove(path)
 
 
 @pytest.mark.skipif(
@@ -602,6 +604,7 @@ def test_tensor_export_cuda(tensor, name):
 
     for s1, s2 in zip(exported.shape, tensor.shape):
         assert s1 == s2
+    os.remove(path)
 
 
 @pytest.mark.skipif(
@@ -627,6 +630,7 @@ def test_tensors_export(tensors, name):
         exported = numpy.load(path)
         exported = exported[exported.files[0]]
         assert numpy.sum(exported.shape) > 1
+        os.remove(path)
 
 
 @flaky(max_runs=2, min_passes=1)

tests/sparseml/sparsification/test_info.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import tempfile
 
 import pytest
@@ -162,3 +163,4 @@ def test_save_load_sparsification_info():
     save_sparsification_info(info, test_path)
     loaded_path = load_sparsification_info(test_path)
     assert info == loaded_path
+    os.remove(test_path)