diff --git a/src/sparseml/onnx/optim/quantization/calibration.py b/src/sparseml/onnx/optim/quantization/calibration.py index 60e123cf9a5..c8f797d5ba0 100644 --- a/src/sparseml/onnx/optim/quantization/calibration.py +++ b/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) diff --git a/src/sparseml/pytorch/torch_to_onnx_exporter.py b/src/sparseml/pytorch/torch_to_onnx_exporter.py index 7d935e97120..cedb9784c52 100644 --- a/src/sparseml/pytorch/torch_to_onnx_exporter.py +++ b/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}") diff --git a/tests/integrations/image_classification/test_image_classification.py b/tests/integrations/image_classification/test_image_classification.py index 2145b259dae..1f5c8fbdb42 100644 --- a/tests/integrations/image_classification/test_image_classification.py +++ b/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): diff --git a/tests/integrations/transformers/test_transformers.py b/tests/integrations/transformers/test_transformers.py index 909b2a0279f..d3a140c590f 100644 --- a/tests/integrations/transformers/test_transformers.py +++ b/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) diff --git a/tests/sparseml/conftest.py b/tests/sparseml/conftest.py index 6d26d4deea0..78ca3a4e54a 100644 --- a/tests/sparseml/conftest.py +++ b/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)}" + ) diff --git a/tests/sparseml/exporters/transforms/test_onnx_transform.py b/tests/sparseml/exporters/transforms/test_onnx_transform.py index 0ac00faee08..de86ad20476 100644 --- a/tests/sparseml/exporters/transforms/test_onnx_transform.py +++ b/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) diff --git a/tests/sparseml/optim/test_analyzer.py b/tests/sparseml/optim/test_analyzer.py index 4c266dbcc03..e18dd88f925 100644 --- a/tests/sparseml/optim/test_analyzer.py +++ b/tests/sparseml/optim/test_analyzer.py @@ -23,7 +23,8 @@ 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) @@ -31,3 +32,4 @@ def test_layer_descs(): 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) diff --git a/tests/sparseml/optim/test_sensitivity.py b/tests/sparseml/optim/test_sensitivity.py index 1a4270837ff..3dae61b8418 100644 --- a/tests/sparseml/optim/test_sensitivity.py +++ b/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() diff --git a/tests/sparseml/pytorch/image_classification/test_export.py b/tests/sparseml/pytorch/image_classification/test_export.py index 414491aa294..4b8c6570988 100644 --- a/tests/sparseml/pytorch/image_classification/test_export.py +++ b/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) diff --git a/tests/sparseml/pytorch/image_classification/utils/test_helpers.py b/tests/sparseml/pytorch/image_classification/utils/test_helpers.py index e1ac7a285ff..bf428a5bfc4 100644 --- a/tests/sparseml/pytorch/image_classification/utils/test_helpers.py +++ b/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) diff --git a/tests/sparseml/pytorch/recipe_template/test_end_to_end.py b/tests/sparseml/pytorch/recipe_template/test_end_to_end.py index 1b35a63e76c..77fe937ad64 100644 --- a/tests/sparseml/pytorch/recipe_template/test_end_to_end.py +++ b/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) diff --git a/tests/sparseml/pytorch/sparsification/test_modifier_thinning.py b/tests/sparseml/pytorch/sparsification/test_modifier_thinning.py index ce3e26ee903..a0a57ea51d4 100644 --- a/tests/sparseml/pytorch/sparsification/test_modifier_thinning.py +++ b/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) diff --git a/tests/sparseml/pytorch/test_torch_to_onnx_exporter.py b/tests/sparseml/pytorch/test_torch_to_onnx_exporter.py index 22f7b212155..66edf77a144 100644 --- a/tests/sparseml/pytorch/test_torch_to_onnx_exporter.py +++ b/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( diff --git a/tests/sparseml/pytorch/utils/test_exporter.py b/tests/sparseml/pytorch/utils/test_exporter.py index cd21d2152b4..3a91a2fb603 100644 --- a/tests/sparseml/pytorch/utils/test_exporter.py +++ b/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(): diff --git a/tests/sparseml/pytorch/utils/test_helpers.py b/tests/sparseml/pytorch/utils/test_helpers.py index 419e5ae30bf..d5fd80b4976 100644 --- a/tests/sparseml/pytorch/utils/test_helpers.py +++ b/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) diff --git a/tests/sparseml/sparsification/test_info.py b/tests/sparseml/sparsification/test_info.py index cf2a0a5071d..b0e7180098c 100644 --- a/tests/sparseml/sparsification/test_info.py +++ b/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)