avoid multiple writes to files on export (#21186)

suo · facebook-github-bot · commit c8dc707fee0e · 2019-05-31T19:14:46.000-07:00
Summary: Pull Request resolved: #21186 ghimport-source-id: 2f62fed Differential Revision: D15581527 Pulled By: suo fbshipit-source-id: b1150cfa47d8df6f217f048c742a5ba9fa7f7935
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -365,6 +365,8 @@ def copy_structure_and_params(m):
 
                 # crack open the zip format to get at the main module code
                 archive = zipfile.ZipFile(buffer)
+                # check that we have no duplicate names
+                self.assertEqual(len(set(archive.namelist())), len(archive.namelist()))
                 main_module = archive.open('archive/code/archive.py')
                 main_module_code = ""
                 for line in main_module:
diff --git a/torch/csrc/jit/export.cpp b/torch/csrc/jit/export.cpp
@@ -607,10 +607,16 @@ void ScriptModuleSerializer::writeLibs(torch::ModelDef* model_def) {
 
   // Write out the files. We still have to do this in converted_classes_ order,
   // to maintain dependency order.
+  std::unordered_set<std::string> written_files;
   for (const auto& item : converted_classes_) {
     const ClassTypePtr& class_type = item.key();
     const std::string filename =
         ImportExportHelpers::qualifierToPath(class_type->qualifier());
+    if (written_files.count(filename)) {
+      continue;
+    }
+    written_files.insert(filename);
+
     const std::string& src = fileToSrc.at(filename).str();
 
     std::ostringstream lib_stream;
