Remove callback from other processes on the same machine

Signed-off-by: Rich Porter <rich.porter@uber.com>

examples/spark/keras/keras_spark_mnist.py

@@ -115,7 +115,8 @@
                                          batch_size=args.batch_size,
                                          epochs=args.epochs,
                                          inmemory_cache_all=True,
-                                         verbose=1)
+                                         verbose=1,
+                                         callbacks=[keras.callbacks.TensorBoard(profile_batch=5)])
 
     keras_model = keras_estimator.fit(train_df).setOutputCols(['label_prob'])
 

horovod/spark/keras/remote.py

@@ -151,6 +151,11 @@ def train(serialized_model, train_rows, val_rows, avg_row_size):
                 # TensorBoard, or other metrics-based callbacks.
                 hvd.callbacks.MetricAverageCallback(),
             ]
+
+            if hvd.local_rank() != 0:
+                # The TB callback appears to conflict across processes on the same machine
+                user_callbacks[:] = [c for c in user_callbacks if not isinstance(c, k.callbacks.TensorBoard)]
+
             callbacks += user_callbacks
 
             # Horovod: save checkpoints only on the first worker to prevent other workers from
@@ -180,12 +185,13 @@ def train(serialized_model, train_rows, val_rows, avg_row_size):
                     for i, c in enumerate(callbacks):
                         if isinstance(c, k.callbacks.TensorBoard):
                             tb_callback = c
+                            print(f"Found TensorBoard callback, updating log_dir to {logs_dir}")
                             tb_callback.log_dir = logs_dir
                             break
                     if tb_callback:
-                        # Rather than a possibly arbitrary order, we always have the TensorBoard
+                        # Rather than a possibly arbitrary order, we always place the TensorBoard
                         # callback right before the SyncCallback
-                        callbacks.remove(i)
+                        callbacks.pop(i)
                     callbacks.append(tb_callback or k.callbacks.TensorBoard(logs_dir))
                     callbacks.append(SyncCallback(run_output_dir, remote_store.sync, k))