I think I finally worked this out.Here is the extra code to save out the model only once and from the 1st node:context = pyspark.BarrierTaskContext.get()
if context.partitionId() == 0: mlflow.keras.log_model(model, "mymodel")
I guess spark_tensorflow_distributor is probably obsolete since there is no update since 2020.Horovod (https://github.com/horovod) seems a better choice on using tensorflow in Databricks with Spark.