add --reload_multifile option to tb launch; use version.parse() in compat lib

Author: Lee Yang

tensorflowonspark/TFSparkNode.py

@@ -137,6 +137,7 @@ def run(fn, tf_args, cluster_meta, tensorboard, log_dir, queues, background):
   """
   def _mapfn(iter):
     import tensorflow as tf
+    from packaging import version
 
     # Note: consuming the input iterator helps Pyspark re-use this worker,
     for i in iter:
@@ -225,7 +226,11 @@ def _mapfn(iter):
         raise Exception("Unable to find 'tensorboard' in: {}".format(search_path))
 
       # launch tensorboard
-      tb_proc = subprocess.Popen([pypath, tb_path, "--logdir=%s" % logdir, "--port=%d" % tb_port], env=os.environ)
+      if version.parse(tf.__version__) >= version.parse('2.0.0'):
+        tb_proc = subprocess.Popen([pypath, tb_path, "--reload_multifile=True", "--logdir=%s" % logdir, "--port=%d" % tb_port], env=os.environ)
+      else:
+        tb_proc = subprocess.Popen([pypath, tb_path, "--logdir=%s" % logdir, "--port=%d" % tb_port], env=os.environ)
+
       tb_pid = tb_proc.pid
 
     # check server to see if this task is being retried (i.e. already reserved)

tensorflowonspark/compat.py

@@ -1,23 +1,22 @@
 # Copyright 2019 Yahoo Inc / Verizon Media
 # Licensed under the terms of the Apache 2.0 license.
 # Please see LICENSE file in the project root for terms.
-"""Helper functions to abstract API changes between TensorFlow versions."""
+"""Helper functions to abstract API changes between TensorFlow versions, intended for end-user TF code."""
 
 import tensorflow as tf
-
-TF_VERSION = tf.__version__
+from packaging import version
 
 
 def export_saved_model(model, export_dir, is_chief=False):
-  if TF_VERSION == '2.0.0':
+  if version.parse(tf.__version__) == version.parse('2.0.0'):
     if is_chief:
       tf.keras.experimental.export_saved_model(model, export_dir)
   else:
     model.save(export_dir, save_format='tf')
 
 
 def disable_auto_shard(options):
-  if TF_VERSION == '2.0.0':
+  if version.parse(tf.__version__) == version.parse('2.0.0'):
     options.experimental_distribute.auto_shard = False
   else:
     options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF