Use ".incomplete" suffix during data generation, rename on complete

PiperOrigin-RevId: 187114891

tensor2tensor/data_generators/generator_utils.py

@@ -147,8 +147,9 @@ def generate_files(generator, output_filenames, max_cases=None):
   if outputs_exist(output_filenames):
     tf.logging.info("Skipping generator because outputs files exist")
     return
+  tmp_filenames = [fname + ".incomplete" for fname in output_filenames]
   num_shards = len(output_filenames)
-  writers = [tf.python_io.TFRecordWriter(fname) for fname in output_filenames]
+  writers = [tf.python_io.TFRecordWriter(fname) for fname in tmp_filenames]
   counter, shard = 0, 0
   for case in generator:
     if case is None:
@@ -165,6 +166,9 @@ def generate_files(generator, output_filenames, max_cases=None):
   for writer in writers:
     writer.close()
 
+  for tmp_name, final_name in zip(tmp_filenames, output_filenames):
+    tf.gfile.Rename(tmp_name, final_name)
+
   tf.logging.info("Generated %s Examples", counter)
 
 

tensor2tensor/data_generators/text_problems.py

@@ -186,10 +186,11 @@ def generate_text_for_vocab(self, data_dir, tmp_dir):
   @property
   def vocab_filename(self):
     if self.vocab_type == VocabType.SUBWORD:
-      return "vocab.%s.%d.%s" % (self.name, self.approx_vocab_size,
+      return "vocab.%s.%d.%s" % (self.dataset_filename(),
+                                 self.approx_vocab_size,
                                  VocabType.SUBWORD)
     else:
-      return "vocab.%s.%s" % (self.name, VocabType.TOKEN)
+      return "vocab.%s.%s" % (self.dataset_filename(), VocabType.TOKEN)
 
   def get_or_create_vocab(self, data_dir, tmp_dir, force_get=False):
     if self.vocab_type == VocabType.CHARACTER: