fix databricks dbfs file path (#4674)

python/orca/src/bigdl/orca/learn/tf/estimator.py

@@ -425,6 +425,12 @@ def to_dataset(data, batch_size, batch_per_thread, validation_data,
     return dataset
 
 
+def save_model_dir(model_dir):
+    if model_dir.startswith("dbfs:/"):
+        model_dir = "/dbfs/" + model_dir[len("dbfs:/"):]
+    return model_dir
+
+
 class TensorFlowEstimator(Estimator):
     def __init__(self, *, inputs, outputs, labels, loss,
                  optimizer, clip_norm, clip_value,
@@ -775,6 +781,8 @@ def shutdown(self):
 
 class KerasEstimator(Estimator):
     def __init__(self, keras_model, metrics, model_dir, optimizer):
+        if model_dir and model_dir.startswith("dbfs:/"):
+            model_dir = save_model_dir(model_dir)
         self.model = KerasModel(keras_model, model_dir)
         self.load_checkpoint = False
         self.metrics = metrics

python/orca/test/bigdl/orca/learn/spark/test_estimator_keras_for_spark.py

@@ -24,6 +24,7 @@
 from bigdl.orca.learn.tf.estimator import Estimator
 from bigdl.dllib.nncontext import *
 from bigdl.orca.learn.utils import convert_predict_rdd_to_dataframe
+from bigdl.orca.learn.tf.estimator import save_model_dir
 
 
 class TestEstimatorForKeras(TestCase):
@@ -681,6 +682,11 @@ def test_estimator_keras_get_model(self):
                 validation_data=df)
         assert est.get_model() is model
 
+    def test_model_path_dbfs_from_keras(self):
+        model_dir = "dbfs:/FileStore/shared_uploads/models"
+        processed_model_dir = save_model_dir(model_dir)
+        assert processed_model_dir == "/dbfs/FileStore/shared_uploads/models"
+
 
 if __name__ == "__main__":
     import pytest