more pep8; py2 vs. py3 compat

Author: leewyang

tensorflowonspark/TFManager.py

@@ -10,6 +10,7 @@
 from multiprocessing.managers import BaseManager
 from multiprocessing import JoinableQueue
 
+
 class TFManager(BaseManager):
   """Python multiprocessing.Manager for distributed, multi-process communication."""
   pass
@@ -20,18 +21,22 @@ class TFManager(BaseManager):
 qdict = {}        # dictionary of queues
 kdict = {}        # dictionary of key-values
 
+
 def _get(key):
   return kdict[key]
 
+
 def _set(key, value):
   kdict[key] = value
 
+
 def _get_queue(qname):
   try:
     return qdict[qname]
   except KeyError:
     return None
 
+
 def start(authkey, queues, mode='local'):
   """Create a new multiprocess.Manager (or return existing one).
 
@@ -53,12 +58,13 @@ def start(authkey, queues, mode='local'):
   TFManager.register('get', callable=lambda key: _get(key))
   TFManager.register('set', callable=lambda key, value: _set(key, value))
   if mode == 'remote':
-    mgr = TFManager(address=('',0), authkey=authkey)
+    mgr = TFManager(address=('', 0), authkey=authkey)
   else:
     mgr = TFManager(authkey=authkey)
   mgr.start()
   return mgr
 
+
 def connect(address, authkey):
   """Connect to a multiprocess.Manager.
 
@@ -75,4 +81,3 @@ def connect(address, authkey):
   m = TFManager(address, authkey=authkey)
   m.connect()
   return m
-

tensorflowonspark/dfutil.py

@@ -109,7 +109,10 @@ def _toTFFeature(name, dtype, row):
       elif dtype in int64_dtypes:
         feature = (name, tf.train.Feature(int64_list=tf.train.Int64List(value=[row[name]])))
       elif dtype in bytes_dtypes:
-        feature = (name, tf.train.Feature(bytes_list=tf.train.BytesList(value=[str(row[name])])))
+        if dtype == 'binary':
+          feature = (name, tf.train.Feature(bytes_list=tf.train.BytesList(value=[bytes(row[name])])))
+        else:
+          feature = (name, tf.train.Feature(bytes_list=tf.train.BytesList(value=[str(row[name]).encode('utf-8')])))
       elif dtype in float_list_dtypes:
         feature = (name, tf.train.Feature(float_list=tf.train.FloatList(value=row[name])))
       elif dtype in int64_list_dtypes:
@@ -181,16 +184,15 @@ def fromTFExample(iter, binary_features=[]):
   """
   # convert from protobuf-like dict to DataFrame-friendly dict
   def _get_value(k, v):
-    # special handling for binary features
-    if k in binary_features:
-      return bytearray(v.bytes_list.value[0])
-
     if v.int64_list.value:
       result = v.int64_list.value
     elif v.float_list.value:
       result = v.float_list.value
-    else:
-      result = v.bytes_list.value
+    else:  # string or bytearray
+      if k in binary_features:
+        return bytearray(v.bytes_list.value[0])
+      else:
+        return v.bytes_list.value[0].decode('utf-8')
 
     if len(result) > 1:         # represent multi-item tensors as python lists
       return list(result)

tensorflowonspark/pipeline.py

@@ -391,7 +391,7 @@ def _fit(self, dataset):
         assert local_args.tfrecord_dir, "Please specify --tfrecord_dir to export DataFrame to TFRecord."
         if self.getInputMapping():
           # if input mapping provided, filter only required columns before exporting
-          dataset = dataset.select(self.getInputMapping().keys())
+          dataset = dataset.select(list(self.getInputMapping()))
         logging.info("Exporting DataFrame {} as TFRecord to: {}".format(dataset.dtypes, local_args.tfrecord_dir))
         dfutil.saveAsTFRecords(dataset, local_args.tfrecord_dir)
         logging.info("Done saving")
@@ -401,7 +401,7 @@ def _fit(self, dataset):
                             local_args.tensorboard, local_args.input_mode, driver_ps_nodes=local_args.driver_ps_nodes)
     if local_args.input_mode == TFCluster.InputMode.SPARK:
       # feed data, using a deterministic order for input columns (lexicographic by key)
-      input_cols = sorted(self.getInputMapping().keys())
+      input_cols = sorted(self.getInputMapping())
       cluster.train(dataset.select(input_cols).rdd, local_args.epochs)
     cluster.shutdown()
 

test/test_TFNode.py

@@ -26,7 +26,7 @@ def test_hdfs_path(self):
 
   def test_datafeed(self):
     """TFNode.DataFeed basic operations"""
-    mgr = TFManager.start('abc', ['input', 'output'], 'local')
+    mgr = TFManager.start('abc'.encode('utf-8'), ['input', 'output'], 'local')
 
     # insert 10 numbers followed by an end-of-feed marker
     q = mgr.get_queue('input')

test/test_dfutil.py

@@ -31,7 +31,7 @@ def test_dfutils(self):
     row1 = ('text string', 1, [2, 3, 4, 5], -1.1, [-2.2, -3.3, -4.4, -5.5], bytearray(b'\xff\xfe\xfd\xfc'))
     rdd = self.sc.parallelize([row1])
     df1 = self.spark.createDataFrame(rdd, ['a', 'b', 'c', 'd', 'e', 'f'])
-    print ("schema: {}".format(df1.schema))
+    print("schema: {}".format(df1.schema))
 
     # save the DataFrame as TFRecords
     dfutil.saveAsTFRecords(df1, self.tfrecord_dir)

test/test_pipeline.py

@@ -2,6 +2,7 @@
 import os
 import shutil
 import test
+import time
 import unittest
 
 from tensorflowonspark import TFCluster, dfutil
@@ -323,6 +324,9 @@ def _get_examples(batch_size):
               ckpt_name = args.model_dir + "/model.ckpt"
               print("Saving checkpoint to: {}".format(ckpt_name))
               saver.save(sess, ckpt_name)
+
+          # wait for rest of cluster to connect
+          time.sleep(30)
         sv.stop()
 
     def _tf_export(args):