tf/estimator.py: only write checkpoint in rank0 (#447)

Author: pang-wu

python/raydp/tf/estimator.py

@@ -188,9 +188,14 @@ def train_func(config):
             if config["evaluate"]:
                 test_history = multi_worker_model.evaluate(eval_tf_dataset, callbacks=callbacks)
                 results.append(test_history)
+
+        # Only save checkpoint from the chief worker to avoid race conditions.
+        # However, we need to call save on all workers to avoid deadlock.
         with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
             multi_worker_model.save(temp_checkpoint_dir, save_format="tf")
-            checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+            checkpoint = None
+            if session.get_world_rank() == 0:
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
 
             session.report({}, checkpoint=checkpoint)