Merge pull request #1016 from pytorch/issue_1012

removing the temp directory with model files after packaging

model-archiver/model_archiver/model_packaging.py

@@ -5,6 +5,7 @@
 
 import logging
 import sys
+import shutil
 from .arg_parser import ArgParser
 from .model_packaging_utils import ModelExportUtils
 from .model_archiver_error import ModelArchiverError
@@ -36,7 +37,7 @@ def package_model(args, manifest):
 
         # Step 2 : Zip 'em all up
         ModelExportUtils.archive(export_file_path, model_name, model_path, manifest, args.archive_format)
-
+        shutil.rmtree(model_path)
         logging.info("Successfully exported model %s to file %s", model_name, export_file_path)
     except ModelArchiverError as e:
         logging.error(e)

model-archiver/model_archiver/tests/unit_tests/test_model_packaging.py

@@ -3,6 +3,11 @@
 from collections import namedtuple
 
 import pytest
+import sys
+from mock import MagicMock
+
+sys.modules['shutil'] = MagicMock()
+sys.modules['shutil.rmtree'] = MagicMock()
 
 from model_archiver.manifest_components.manifest import RuntimeType
 from model_archiver.model_packaging import generate_model_archive, package_model