Adding metrics and return results.

Change-Id: Icc1efe26b51f7e3780d9724f488af1f730d0d15b

Author: mbernico

requirements.txt

@@ -9,4 +9,4 @@ pylint >= 2.5.3
 fire >= 0.3.1
 jupyter >= 1.0.0
 tensorflow >= 2.2.0
-gcsfs >= 0.6.2
+pyarrow < 0.17

setup.py

@@ -29,7 +29,6 @@
     "pylint >= 2.5.3",
     "fire >= 0.3.1",
     "tensorflow >= 2.2.0",
-    "gcsfs >= 0.6.2",
     "pyarrow < 0.17",
 ]
 

tfrutil/accessor.py

@@ -22,8 +22,10 @@
 """
 from typing import Any, Dict, Optional, Union
 import pandas as pd
+from IPython.core import display
 
 from tfrutil import client
+from tfrutil import constants
 
 
 @pd.api.extensions.register_dataframe_accessor('tensorflow')
@@ -43,7 +45,7 @@ def to_tfr(
       dataflow_options: Union[Dict[str, Any], None] = None,
       job_label: str = 'to-tfr',
       compression: Optional[str] = 'gzip',
-      num_shards: int = 0):
+      num_shards: int = 0) -> Dict[str, Any]:
     """TFRUtil Pandas Accessor.
 
     TFRUtil provides an easy interface to create image-based tensorflow records
@@ -68,8 +70,14 @@ def to_tfr(
       compression: Can be 'gzip' or None for no compression.
       num_shards: Number of shards to divide the TFRecords into. Default is
           0 = no sharding.
+    Returns:
+      job_results: A dictionary of job results.
     """
-    client.create_tfrecords(
+    display.display(
+        display.HTML(
+            '<b>Logging output to /tmp/{} </b>'.format(constants.LOGFILE)))
+
+    r = client.create_tfrecords(
         self._df,
         output_dir=output_dir,
         runner=runner,
@@ -79,4 +87,4 @@ def to_tfr(
         job_label=job_label,
         compression=compression,
         num_shards=num_shards)
-    #TODO (mikebernico) Add notebook output for user.
+    return r

tfrutil/accessor_test.py

@@ -36,8 +36,9 @@ def setUp(self):
   def test_accessor(self):
     """Tests pandas accessor."""
 
-    self.assertIsNone(self.test_df.tensorflow.to_tfr(
-        runner='DirectRunner', output_dir=self.output_dir))
+    r = self.test_df.tensorflow.to_tfr(
+        runner='DirectRunner', output_dir=self.output_dir)
+    self.assertTrue('metrics' in r)
 
 if __name__ == '__main__':
   unittest.main()

tfrutil/beam_image.py

@@ -20,6 +20,7 @@
 from typing import Any, Dict, Generator, Tuple
 
 import apache_beam as beam
+from apache_beam.metrics import Metrics
 import tensorflow as tf
 from PIL import Image
 
@@ -76,11 +77,14 @@ def __init__(self, image_key: str):
     """Constructor."""
     super().__init__()
     self.image_key = image_key
+    self.image_good_counter = Metrics.counter(self.__class__, 'image_good')
+    self.image_bad_counter = Metrics.counter(self.__class__, 'image_bad')
+
 
   # pylint: disable=unused-argument
   def process(
       self,
-      element: Dict,
+      element: Dict[str, Any],
       *args: Tuple[Any, ...],
       **kwargs: Dict) -> Generator[Dict[str, Any], None, None]:
     """Loads image and creates image features.
@@ -96,11 +100,13 @@ def process(
       d['image'] = encode(image)
       d['image_width'], d['image_height'] = image.size
       d['image_channels'] = mode_to_channel(image.mode)
+      self.image_good_counter.inc()
 
     # pylint: disable=broad-except
     except Exception as e:
       logging.warning('Could not load image: %s', image_uri)
       logging.error('Exception was: %s', str(e))
+      self.image_bad_counter.inc()
 
     element.update(d)
     yield element

tfrutil/beam_pipeline.py

@@ -19,10 +19,11 @@
 This file implements the full beam pipeline for TFRUtil.
 """
 
+from typing import Any, Dict, Generator, Union
+
 import functools
 import logging
 import os
-from typing import Any, Dict, Union
 
 import apache_beam as beam
 import pandas as pd
@@ -159,6 +160,34 @@ def _preprocessing_fn(inputs, integer_label: bool = False):
   return outputs
 
 
+# pylint: disable=abstract-method
+
+class ToCSVRows(beam.DoFn):
+  """Adds image to PCollection."""
+
+  def __init__(self):
+    """Constructor."""
+    super().__init__()
+    self.row_count = beam.metrics.Metrics.counter(self.__class__, 'row_count')
+
+
+  # pylint: disable=unused-argument
+  # pylint: disable=arguments-differ
+  def process(
+      self,
+      element: Dict[str, Any]
+      ) -> Generator[Dict[str, Any], None, None]:
+    """Loads image and creates image features.
+
+    This DoFn extracts an image being stored on local disk or GCS and
+    yields a base64 encoded image, the image height, image width, and channels.
+    """
+    element = ','.join([str(item) for item in element])
+    self.row_count.inc()
+    yield element
+
+
+
 # pylint: disable=too-many-arguments
 # pylint: disable=too-many-locals
 def build_pipeline(
@@ -210,15 +239,15 @@ def build_pipeline(
                                     constants.IMAGE_CSV_METADATA.schema)
 
     extract_images_fn = beam_image.ExtractImagesDoFn(constants.IMAGE_URI_KEY)
+    flatten_rows = ToCSVRows()
 
     # Each element in the image_csv_data PCollection will be a dict
     # including the image_csv_columns and the image features created from
     # extract_images_fn.
     image_csv_data = (
         p
         | 'ReadFromDataFrame' >> beam.Create(df.values.tolist())
-        | 'ToCSVRows' >> beam.Map(
-            lambda x: ','.join([str(item) for item in x]))
+        | 'ToCSVRows' >> beam.ParDo(flatten_rows)
         | 'DecodeCSV' >> beam.Map(converter.decode)
         | 'ReadImage' >> beam.ParDo(extract_images_fn)
     )

tfrutil/client.py

@@ -23,10 +23,11 @@
 import os
 from typing import Any, Dict, Union, Optional, Sequence
 
+import apache_beam as beam
 import pandas as pd
 import tensorflow as tf
 
-# from tfrutil import common
+from tfrutil import common
 from tfrutil import constants
 from tfrutil import beam_pipeline
 
@@ -118,6 +119,40 @@ def to_dataframe(
 
   return df
 
+
+def _get_beam_metric(
+    metric_filter: beam.metrics.MetricsFilter,
+    result: beam.runners.runner.PipelineResult,
+    metric_type: str = 'counters') -> Optional[int]:
+  """Queries a beam pipeline result for a specificed metric.
+
+  Args:
+    metric_filter: an instance of apache_beam.metrics.MetricsFilter()
+    metric_type: A metric type (counters, distributions, etc.)
+
+  Returns:
+    Counter value or None
+  """
+  query_result = result.metrics().query(metric_filter)
+  result_val = None
+  if query_result[metric_type]:
+    result_val = query_result[metric_type][0].result
+  return result_val
+
+
+def _configure_logging(logfile):
+  """Configures logging options."""
+  # Remove default handlers that TF set for us.
+  logger = logging.getLogger('')
+  logger.handlers = []
+  handler = logging.FileHandler(logfile)
+  logger.addHandler(handler)
+  logger.setLevel(constants.LOGLEVEL)
+  # This disables annoying Tensorflow and TFX info/warning messages on console.
+  tf_logger = logging.getLogger('tensorflow')
+  tf_logger.handlers = []
+  tf_logger.addHandler(handler)
+
 # pylint: disable=too-many-arguments
 # pylint: disable=too-many-locals
 
@@ -132,7 +167,7 @@ def create_tfrecords(
     dataflow_options: Optional[Dict[str, Any]] = None,
     job_label: str = 'create-tfrecords',
     compression: Optional[str] = 'gzip',
-    num_shards: int = 0):
+    num_shards: int = 0) -> Dict[str, Any]:
   """Generates TFRecord files from given input data.
 
   TFRUtil provides an easy interface to create image-based tensorflow records
@@ -161,18 +196,21 @@ def create_tfrecords(
     num_shards: Number of shards to divide the TFRecords into. Default is
         0 = no sharding.
 
+  Returns:
+    job_results: Dict
+      job_id: DataFlow Job ID or 'DirectRunner'
+      metrics: (optional) Beam metrics. Only used for DirectRunner
+      dataflow_url: (optional) Job URL for DataFlowRunner
   """
 
   df = to_dataframe(input_data, header, names)
 
   _validate_data(df)
   _validate_runner(df, runner, project, region)
-  #os.makedirs(output_dir, exist_ok=True)
-  #TODO (mikebernico) this doesn't work with GCS locations...
+
   logfile = os.path.join('/tmp', constants.LOGFILE)
-  logging.basicConfig(filename=logfile, level=constants.LOGLEVEL)
-  # This disables annoying Tensorflow and TFX info/warning messages.
-  logging.getLogger('tensorflow').setLevel(logging.ERROR)
+  _configure_logging(logfile)
+
 
   integer_label = pd.api.types.is_integer_dtype(df[constants.LABEL_KEY])
   p = beam_pipeline.build_pipeline(
@@ -187,15 +225,55 @@ def create_tfrecords(
       dataflow_options=dataflow_options,
       integer_label=integer_label)
 
-  # TODO(mikbernico) Handle this async for the DataFlow case.
   result = p.run()
-  result.wait_until_finish()
-  # TODO(mikebernico) Add metrics here.
+
+  if runner == 'DirectRunner':
+    logging.info("Using DirectRunner. TFRUtil will block until job completes.")
+    result.wait_until_finish()
+
+    row_count_filter = beam.metrics.MetricsFilter().with_name('row_count')
+    good_image_filter = beam.metrics.MetricsFilter().with_name('image_good')
+    bad_image_filter = beam.metrics.MetricsFilter().with_name('image_bad')
+
+    row_count = _get_beam_metric(row_count_filter, result)
+    good_image_count = _get_beam_metric(good_image_filter, result)
+    bad_image_count = _get_beam_metric(bad_image_filter, result)
+
+    # TODO(mikebernico): Profile metric impact with larger dataset.
+    metrics = {
+        'rows': row_count,
+        'good_images': good_image_count,
+        'bad_images': bad_image_count,
+    }
+
+    job_result = {
+        'job_id': 'DirectRunner',
+        'metrics': metrics
+    }
+    logging.info("Job Complete.")
+
+  else:
+    logging.info("Using DataFlow Runner.")
+    # Construct DataFlow URL
+
+    job_id = result.job_id()
+
+    url = (
+        constants.CONSOLE_DATAFLOW_URI +
+        region +
+        '/' +
+        job_id +
+        '?project=' +
+        project)
+    job_result = {
+        'job_id': job_id,
+        'dataflow_url': url
+    }
+
   logging.shutdown()
 
-  # FIXME: Issue where GCSFS is not picking up the `logfile` even if it exists.
-  if os.path.exists(logfile):
-    pass
-    # common.copy_to_gcs(logfile,
-    #                    os.path.join(output_dir, constants.LOGFILE),
-    #                    recursive=False)
+  if runner == 'DataFlowRunner':
+    # if this is a dataflow job, copy the logfile to gcs
+    common.copy_logfile_to_gcs(logfile, output_dir)
+
+  return job_result

tfrutil/client_test.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 """Tests for client."""
-
+import os
 from typing import List
 
 import csv
@@ -35,14 +35,43 @@ class ClientTest(unittest.TestCase):
 
   def setUp(self):
     self.test_df = test_utils.get_test_df()
+    self.test_region = 'us-central1'
+    self.test_project = 'foo'
 
-  def test_create_tfrecords(self):
-    """Tests `create_tfrecords` valid case."""
-
-    self.assertIsNone(client.create_tfrecords(
+  @mock.patch('tfrutil.client.beam_pipeline')
+  def test_create_tfrecords_direct_runner(self, mock_beam):
+    """Tests `create_tfrecords` Direct case."""
+    mock_beam.build_pipeline().run().wait_until_finished.return_value = {
+        'rows':6}
+    r = client.create_tfrecords(
         self.test_df,
         runner='DirectRunner',
-        output_dir='/tmp/train'))
+        output_dir='/tmp/direct_runner')
+    self.assertTrue('metrics' in r)
+
+  @mock.patch('tfrutil.client.beam_pipeline')
+  def test_create_tfrecords_dataflow_runner(self, mock_beam):
+    """Tests `create_tfrecords` DataFlow case."""
+    mock_beam.build_pipeline().run().job_id.return_value = 'foo_id'
+
+    df2 = self.test_df.copy()
+    df2[constants.IMAGE_URI_KEY] = 'gs://' + df2[constants.IMAGE_URI_KEY]
+
+    outdir = '/tmp/dataflow_runner'
+
+    expected = {
+        'job_id': 'foo_id',
+        'dataflow_url': 'https://console.cloud.google.com/dataflow/jobs/' +
+                        'us-central1/foo_id?project=foo'}
+
+    os.makedirs(outdir, exist_ok=True)
+    r = client.create_tfrecords(
+        df2,
+        runner='DataFlowRunner',
+        output_dir=outdir,
+        region=self.test_region,
+        project=self.test_project)
+    self.assertEqual(r, expected)
 
 
 # pylint: disable=protected-access