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

* initial commit

* Update src/deepsparse/license.py

* limit to 150mb

* ready to review

* initial commit

* Update src/deepsparse/license.py

* limit to 150mb

* ready to review

* fix the erronous Makefile

* perhaps fixed GHA

* take into consideration that GHA creates four files

Makefile

@@ -53,12 +53,12 @@ artifacts:
 # run tests for the repo
 test:
 	@echo "Running python tests";
-	@SPARSEZOO_TEST_MODE="true" @NM_DISABLE_ANALYTICS="true" pytest tests/ --ignore integrations $(PYTEST_ARGS);
+	@SPARSEZOO_TEST_MODE="true" NM_DISABLE_ANALYTICS="true" pytest tests/ --ignore integrations $(PYTEST_ARGS);
 
 # run integrations tests for the repo
 test_integrations:
 	@echo "Running package integrations tests";
-	@SPARSEZOO_TEST_MODE="true" @NM_DISABLE_ANALYTICS="true" pytest integrations/ --ignore tests $(PYTEST_ARGS);
+	@SPARSEZOO_TEST_MODE="true" NM_DISABLE_ANALYTICS="true" pytest integrations/ --ignore tests $(PYTEST_ARGS);
 
 # create docs
 docs:

src/deepsparse/license.py

@@ -53,7 +53,7 @@
 def add_deepsparse_license(token_or_path):
     candidate_license_file_path = token_or_path
     if not os.path.exists(token_or_path):
-        # write raw token to temp file for validadation
+        # write raw token to temp file for validation
         candidate_license_tempfile = NamedTemporaryFile()
         candidate_license_file_path = candidate_license_tempfile.name
         with open(candidate_license_file_path, "w") as token_file:

src/deepsparse/server/cli.py

@@ -228,6 +228,7 @@ def main(
             loggers={},
         )
 
+        # saving yaml config to temporary directory
         with TemporaryDirectory() as tmp_dir:
             config_path = os.path.join(tmp_dir, "server-config.yaml")
             with open(config_path, "w") as fp:

src/deepsparse/transformers/helpers.py

@@ -135,7 +135,7 @@ def overwrite_transformer_onnx_model_inputs(
     path: str,
     batch_size: int = 1,
     max_length: int = 128,
-    output_path: Optional[str] = None,
+    inplace: bool = True,
 ) -> Tuple[Optional[str], List[str], Optional[NamedTemporaryFile]]:
     """
     Overrides an ONNX model's inputs to have the given batch size and sequence lengths.
@@ -145,15 +145,17 @@ def overwrite_transformer_onnx_model_inputs(
     :param path: path to the ONNX model to override
     :param batch_size: batch size to set
     :param max_length: max sequence length to set
-    :param output_path: if provided, the model will be saved to the given path,
-        otherwise, the model will be saved to a named temporary file that will
-        be deleted after the program exits
-    :return: if no output path, a tuple of the saved path to the model, list of
-        model input names, and reference to the tempfile object will be returned
-        otherwise, only the model input names will be returned
+    :param inplace: if True, the model will be modified in place (its inputs will
+        be overwritten). Else, a copy of that model, with overwritten inputs,
+        will be saved to a temporary file
+    :return: tuple of (path to the overwritten model, list of input names that were
+        overwritten, and a temporary file containing the overwritten model if
+        `inplace=False`, else None)
     """
     # overwrite input shapes
-    model = onnx.load(path)
+    # if > 2Gb model is to be modified in-place, operate
+    # exclusively on the model graph
+    model = onnx.load(path, load_external_data=not inplace)
     initializer_input_names = set([node.name for node in model.graph.initializer])
     external_inputs = [
         inp for inp in model.graph.input if inp.name not in initializer_input_names
@@ -165,14 +167,20 @@ def overwrite_transformer_onnx_model_inputs(
         input_names.append(external_input.name)
 
     # Save modified model
-    if output_path is None:
-        tmp_file = NamedTemporaryFile()  # file will be deleted after program exit
+    if inplace:
+        _LOGGER.info(
+            f"Overwriting in-place the input shapes of the transformer model at {path}"
+        )
+        save_onnx(model, path)
+        return path, input_names, None
+    else:
+        tmp_file = NamedTemporaryFile()
+        _LOGGER.info(
+            f"Saving a copy of the transformer model: {path} "
+            f"with overwritten input shapes to {tmp_file.name}"
+        )
         save_onnx(model, tmp_file.name)
         return tmp_file.name, input_names, tmp_file
-    else:
-        save_onnx(model, output_path)
-
-        return input_names
 
 
 def _get_file_parent(file_path: str) -> str:

src/deepsparse/utils/onnx.py

@@ -53,18 +53,22 @@
 
 
 @contextlib.contextmanager
-def save_onnx_to_temp_files(model: Model, with_external_data=True) -> str:
+def save_onnx_to_temp_files(model: onnx.ModelProto, with_external_data=False) -> str:
     """
     Save model to a temporary file. Works for models with external data.
+
     :param model: The onnx model to save to temporary directory
     :param with_external_data: Whether to save external data to a separate file
     """
     shaped_model = tempfile.NamedTemporaryFile(mode="w", delete=False)
+    _LOGGER.info(f"Saving model to temporary directory: {tempfile.tempdir}")
+
     if with_external_data:
         external_data = os.path.join(
             tempfile.tempdir, next(tempfile._get_candidate_names())
         )
         has_external_data = save_onnx(model, shaped_model.name, external_data)
+        _LOGGER.info(f"Saving external data to temporary directory: {external_data}")
     else:
         has_external_data = save_onnx(model, shaped_model.name)
     try:
@@ -194,17 +198,29 @@ def generate_random_inputs(
     return input_data_list
 
 
+@contextlib.contextmanager
 def override_onnx_batch_size(
-    onnx_filepath: str, batch_size: int, inplace: bool = False
+    onnx_filepath: str,
+    batch_size: int,
+    inplace: bool = True,
 ) -> str:
     """
     Rewrite batch sizes of ONNX model, saving the modified model and returning its path
-    :param onnx_filepath: File path to ONNX model
+
+    :param onnx_filepath: File path to ONNX model. If the graph is to be
+        modified in-place, only the model graph will be loaded and modified.
+        Otherwise, the entire model will be loaded and modified, so that
+        external data are saved along the model graph.
     :param batch_size: Override for the batch size dimension
-    :param inplace: If True, overwrite the original model file
-    :return: File path to modified ONNX model
+    :param inplace: If True, overwrite the original model file.
+        Else, save the modified model to a temporary file.
+    :return: File path to modified ONNX model.
+        If inplace is True,
+        the modified model will be saved to the same path as the original
+        model. Else the modified model will be saved to a
+        temporary file.
     """
-    model = onnx.load(onnx_filepath, load_external_data=False)
+    model = onnx.load(onnx_filepath, load_external_data=not inplace)
     all_inputs = model.graph.input
     initializer_input_names = [node.name for node in model.graph.initializer]
     external_inputs = [
@@ -213,32 +229,41 @@ def override_onnx_batch_size(
     for external_input in external_inputs:
         external_input.type.tensor_type.shape.dim[0].dim_value = batch_size
 
-    # Save modified model, this will be cleaned up when context is exited
     if inplace:
-        onnx.save(model, onnx_filepath)
-        return onnx_filepath
+        _LOGGER.info(
+            f"Overwriting in-place the batch size of the model at {onnx_filepath}"
+        )
+        save_onnx(model, onnx_filepath)
+        yield onnx_filepath
     else:
-        # Save modified model, this will be cleaned up when context is exited
-        return save_onnx_to_temp_files(model, with_external_data=False)
+        return save_onnx_to_temp_files(model, with_external_data=not inplace)
 
 
 def override_onnx_input_shapes(
     onnx_filepath: str,
     input_shapes: Union[List[int], List[List[int]]],
-    inplace: bool = False,
+    inplace: bool = True,
 ) -> str:
     """
     Rewrite input shapes of ONNX model, saving the modified model and returning its path
-    :param onnx_filepath: File path to ONNX model
+
+    :param onnx_filepath: File path to ONNX model. If the graph is to be
+        modified in-place, only the model graph will be loaded and modified.
+        Otherwise, the entire model will be loaded and modified, so that
+        external data are saved along the model graph.
     :param input_shapes: Override for model's input shapes
     :param inplace: If True, overwrite the original model file
-    :return: File path to modified ONNX model
+    :return: File path to modified ONNX model.
+        If inplace is True,
+        the modified model will be saved to the same path as the original
+        model. Else the modified model will be saved to a
+        temporary file.
     """
 
     if input_shapes is None:
         return onnx_filepath
 
-    model = onnx.load(onnx_filepath, load_external_data=False)
+    model = onnx.load(onnx_filepath, load_external_data=not inplace)
     all_inputs = model.graph.input
     initializer_input_names = [node.name for node in model.graph.initializer]
     external_inputs = [
@@ -273,13 +298,18 @@ def override_onnx_input_shapes(
         for dim_idx, dim in enumerate(external_input.type.tensor_type.shape.dim):
             dim.dim_value = input_shapes[input_idx][dim_idx]
 
-    # Save modified model, this will be cleaned up when context is exited
     if inplace:
+        _LOGGER.info(
+            "Overwriting in-place the input shapes of the model " f"at {onnx_filepath}"
+        )
         onnx.save(model, onnx_filepath)
         return onnx_filepath
     else:
-        # Save modified model, this will be cleaned up when context is exited
-        return save_onnx_to_temp_files(model, with_external_data=False)
+        _LOGGER.info(
+            f"Saving the input shapes of the model at {onnx_filepath} "
+            f"to a temporary file"
+        )
+        return save_onnx_to_temp_files(model, with_external_data=not inplace)
 
 
 def truncate_onnx_model(
@@ -358,6 +388,7 @@ def truncate_onnx_model(
             output.type.tensor_type.shape.Clear()
 
     # save and check model
+    _LOGGER.debug(f"Saving truncated model to {output_filepath}")
     save_onnx(extracted_model, output_filepath, "external_data")
     validate_onnx(output_filepath)
 

src/deepsparse/yolo/utils/utils.py

@@ -29,6 +29,7 @@
 import yaml
 
 import torch
+from deepsparse.utils.onnx import save_onnx_to_temp_files
 from deepsparse.yolo.schemas import YOLOOutput
 from sparsezoo.utils import save_onnx
 
@@ -341,7 +342,7 @@ def get_onnx_expected_image_shape(onnx_model: onnx.ModelProto) -> Tuple[int, ...
 
 
 def modify_yolo_onnx_input_shape(
-    model_path: str, image_shape: Tuple[int, int]
+    model_path: str, image_shape: Tuple[int, int], inplace: bool = True
 ) -> Tuple[str, Optional[NamedTemporaryFile]]:
     """
     Creates a new YOLO ONNX model from the given path that accepts the given input
@@ -350,13 +351,17 @@ def modify_yolo_onnx_input_shape(
 
     :param model_path: file path to YOLO ONNX model
     :param image_shape: 2-tuple of the image shape to resize this yolo model to
-    :return: filepath to an onnx model reshaped to the given input shape will be the
-        original path if the shape is the same.  Additionally returns the
-        NamedTemporaryFile for managing the scope of the object for file deletion
+    :param inplace: if True, modifies the given model_path in-place, otherwise
+        saves the modified model to a temporary file
+    :return: filepath to an onnx model reshaped to the given input shape.
+        If inplace is True,
+        the modified model will be saved to the same path as the original
+        model. Else the modified model will be saved to a
+        temporary file.
     """
     has_postprocessing = yolo_onnx_has_postprocessing(model_path)
 
-    model = onnx.load(model_path)
+    model = onnx.load(model_path, load_external_data=not inplace)
     model_input = model.graph.input[0]
 
     initial_x, initial_y = get_onnx_expected_image_shape(model)
@@ -399,10 +404,18 @@ def modify_yolo_onnx_input_shape(
         )
         set_tensor_dim_shape(model.graph.output[0], 1, num_predictions)
 
-    tmp_file = NamedTemporaryFile()  # file will be deleted after program exit
-    save_onnx(model, tmp_file.name)
-
-    return tmp_file.name, tmp_file
+    if inplace:
+        _LOGGER.info(
+            "Overwriting in-place the ONNX model "
+            f"at {model_path} with the new input shape"
+        )
+        save_onnx(model, model_path)
+        return model_path
+    else:
+        _LOGGER.info(
+            "Saving the ONNX model with the " "new input shape to a temporary file"
+        )
+        return save_onnx_to_temp_files(model, with_external_data=not inplace)
 
 
 def get_tensor_dim_shape(tensor: onnx.TensorProto, dim: int) -> int:

tests/conftest.py

@@ -13,13 +13,22 @@
 # limitations under the License.
 
 import os
+import tempfile
 from subprocess import Popen
 from typing import List
 
 import pytest
 from tests.helpers import delete_file
 
 
+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
 def cleanup():
     filenames: List[str] = []
@@ -50,3 +59,34 @@ def cleanup():
         )
         for proc in processes:
             proc.terminate()
+
+
+@pytest.fixture(scope="session", autouse=True)
+def check_for_created_files():
+    start_files_root = _get_files(directory=r".")
+    start_files_temp = _get_files(directory=tempfile.gettempdir())
+    yield
+    end_files_root = _get_files(directory=r".")
+    end_files_temp = _get_files(directory=tempfile.gettempdir())
+
+    max_allowed_number_created_files = 4
+    # GHA needs to create following files:
+    # pyproject.toml, CONTRIBUTING.md, LICENSE, setup.cfg
+    assert len(start_files_root) + max_allowed_number_created_files >= 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 = 150
+    size_of_temp_files_bytes = sum(
+        os.path.getsize(path) for path in set(end_files_temp) - set(start_files_temp)
+    )
+    size_of_temp_files_megabytes = size_of_temp_files_bytes / 1024 / 1024
+    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/deepsparse/loggers/metric_functions/helpers/test_config_generation.py

@@ -14,6 +14,7 @@
 
 
 import os
+import shutil
 
 import yaml
 
@@ -155,6 +156,8 @@ def test_data_logging_config_from_predefined(
         with open(os.path.join(tmp_path, "data_logging_config.yaml"), "r") as stream:
             string_result_saved = yaml.safe_load(stream)
         assert string_result_saved == yaml.safe_load(expected_result)
+        return
+    shutil.rmtree(tmp_path, ignore_errors=True)
 
 
 result_1 = """loggers:

tests/deepsparse/loggers/test_prometheus_logger.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+import shutil
+
 import requests
 
 import pytest
@@ -119,6 +121,7 @@ def test_prometheus_logger(
     count_request_text = float(text_log_lines[98].split(" ")[1])
 
     assert count_request_request == count_request_text == no_iterations
+    shutil.rmtree(tmp_path)
 
 
 @pytest.mark.parametrize(