tf/estimator.py: only write checkpoint in rank0

Author: pang-wu

python/raydp/tf/estimator.py

@@ -188,11 +188,15 @@ def train_func(config):
             if config["evaluate"]:
                 test_history = multi_worker_model.evaluate(eval_tf_dataset, callbacks=callbacks)
                 results.append(test_history)
-        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
-            multi_worker_model.save(temp_checkpoint_dir, save_format="tf")
-            checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+        
+        # Only save checkpoint from the chief worker to avoid race conditions
+        checkpoint = None
+        if session.get_world_rank() == 0:
+            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                multi_worker_model.save(temp_checkpoint_dir, save_format="tf")
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
 
-            session.report({}, checkpoint=checkpoint)
+        session.report({}, checkpoint=checkpoint)
 
     def fit(self,
             train_ds: Dataset,
@@ -290,6 +294,21 @@ def fit_on_spark(self,
 
     def get_model(self) -> Any:
         assert self._trainer, "Trainer has not been created"
-        return TensorflowCheckpoint.from_saved_model(
-                self._results.checkpoint.to_directory()
-            ).get_model()
+        checkpoint_dir = self._results.checkpoint.to_directory()
+        
+        try:
+            # Try standard loading
+            return TensorflowCheckpoint.from_saved_model(checkpoint_dir).get_model()
+        except RuntimeError as e:
+            if "Fingerprint" in str(e) or "fingerprint.pb" in str(e):
+                # Fallback: Load directly with tf.keras.models.load_model
+                # This bypasses Ray Train's checkpoint wrapper and fingerprint validation
+                import warnings
+                warnings.warn(
+                    f"Encountered fingerprint error when loading checkpoint: {e}. "
+                    "Falling back to direct tf.keras.models.load_model. "
+                    "This may indicate a TensorFlow version incompatibility."
+                )
+                return tf.keras.models.load_model(checkpoint_dir)
+            else:
+                raise