Add comments, exit async process after waiting for flagfile for too long and...

Add comments, exit async process after waiting for flagfile for too long and make directory for data_dir in case it does not exist.

official/recommendation/data_async_generation.py

@@ -441,14 +441,21 @@ def _generation_loop(num_workers,           # type: int
 
 
 def _parse_flagfile():
-  """Fill flags with flagfile."""
+  """Fill flags with flagfile written by the main process."""
   flagfile = os.path.join(flags.FLAGS.data_dir,
                           rconst.FLAGFILE)
   tf.logging.info("Waiting for flagfile to appear at {}..."
                   .format(flagfile))
+  start_time = time.time()
   while not tf.gfile.Exists(flagfile):
+    if time.time() - start_time > rconst.TIMEOUT_SECONDS:
+      log_msg("Waited more than {} seconds. Concluding that this "
+              "process is orphaned and exiting gracefully."
+              .format(rconst.TIMEOUT_SECONDS))
+      sys.exit()
     time.sleep(1)
   tf.logging.info("flagfile found.")
+  # This overrides FLAGS with flags from flagfile.
   flags.FLAGS([__file__, "--flagfile", flagfile])
 
 

official/recommendation/data_preprocessing.py

@@ -450,6 +450,7 @@ def instantiate_pipeline(dataset, data_dir, batch_size, eval_batch_size,
   }
   if ncf_dataset.deterministic:
     flags_["seed"] = stat_utils.random_int32()
+  tf.gfile.MakeDirs(flags.FLAGS.data_dir)
   # We write to a temp file then atomically rename it to the final file,
   # because writing directly to the final file can cause the data generation
   # async process to read a partially written JSON file.