Removed tfrutil_path option

Change-Id: I3affb81cb5fe4d8af837f570cbeda8a4f7292d64

Author: mbernico

tfrutil/accessor.py

@@ -40,7 +40,6 @@ def to_tfr(
       runner: str = 'DirectRunner',
       project: Optional[str] = None,
       region: Optional[str] = None,
-      tfrutil_path: Optional[str] = None,
       dataflow_options: Union[Dict[str, Any], None] = None,
       job_label: str = 'to-tfr',
       compression: Optional[str] = 'gzip',
@@ -64,7 +63,6 @@ def to_tfr(
       runner: Beam runner. Can be DirectRunner or  DataFlowRunner.
       project: GCP project name (Required if DataFlowRunner).
       region: GCP region name (Required if DataFlowRunner).
-      tfrutil_path: Path to tfrutil source (Required if DataFlowRunner).
       dataflow_options: Optional dictionary containing DataFlow options.
       job_label: User supplied description for the beam job name.
       compression: Can be 'gzip' or None for no compression.
@@ -77,7 +75,6 @@ def to_tfr(
         runner=runner,
         project=project,
         region=region,
-        tfrutil_path=tfrutil_path,
         dataflow_options=dataflow_options,
         job_label=job_label,
         compression=compression,

tfrutil/beam_pipeline.py

@@ -65,7 +65,6 @@ def _get_pipeline_options(
     job_dir: str,
     project: str,
     region: str,
-    tfrutil_path: str,
     dataflow_options: Union[Dict[str, Any], None]
     ) -> beam.pipeline.PipelineOptions:
   """Returns Beam pipeline options."""
@@ -84,8 +83,8 @@ def _get_pipeline_options(
     options_dict['project'] = project
   if region:
     options_dict['region'] = region
-  if tfrutil_path:
-    options_dict['setup_file'] = os.path.join(tfrutil_path, 'setup.py')
+  if runner == 'DataFlowRunner':
+    options_dict['setup_file'] = os.path.join('..', 'setup.py')
   if dataflow_options:
     options_dict.update(dataflow_options)
 
@@ -158,7 +157,6 @@ def build_pipeline(
     runner: str,
     project: str,
     region: str,
-    tfrutil_path: str,
     output_dir: str,
     compression: str,
     num_shards: int,
@@ -172,7 +170,6 @@ def build_pipeline(
     runner: Beam Runner: (e.g. DataFlowRunner, DirectRunner).
     project: GCP project ID (if DataFlowRunner)
     region: GCP compute region (if DataFlowRunner)
-    tfrutil_path: Path for TFRUtil source (required for DataFlowRunner)
     output_dir: GCS or Local Path for output.
     compression: gzip or None.
     num_shards: Number of shards.
@@ -193,7 +190,6 @@ def build_pipeline(
       job_dir,
       project,
       region,
-      tfrutil_path,
       dataflow_options)
 
   #with beam.Pipeline(runner, options=options) as p:

tfrutil/client.py

@@ -53,8 +53,7 @@ def _validate_runner(
     df: pd.DataFrame,
     runner: str,
     project: str,
-    region: str,
-    tfrutil_path: str):
+    region: str):
   """Validates an appropriate beam runner is chosen."""
   if runner not in ['DataFlowRunner', 'DirectRunner']:
     raise AttributeError('Runner {} is not supported.'.format(runner))
@@ -65,11 +64,10 @@ def _validate_runner(
     raise AttributeError('DataFlowRunner requires GCS image locations.')
 
   if (runner == 'DataFlowRunner') & (
-      any(not v for v in [project, region, tfrutil_path])):
-    raise AttributeError('DataFlowRunner requires project region and '
-                         'tfrutil_path to be set however project is {} '
-                         'region is {} and tfrutil_path is {}'.format(
-                             project, region, tfrutil_path))
+      any(not v for v in [project, region])):
+    raise AttributeError('DataFlowRunner requires project region and region '
+                         'project is {} and region is {}'.format(
+                             project, region))
 
 # def read_image_directory(dirpath) -> pd.DataFrame:
 #   """Reads image data from a directory into a Pandas DataFrame."""
@@ -131,7 +129,6 @@ def create_tfrecords(
     runner: str = 'DirectRunner',
     project: Optional[str] = None,
     region: Optional[str] = None,
-    tfrutil_path: Optional[str] = None,
     dataflow_options: Optional[Dict[str, Any]] = None,
     job_label: str = 'create-tfrecords',
     compression: Optional[str] = 'gzip',
@@ -158,7 +155,6 @@ def create_tfrecords(
     runner: Beam runner. Can be 'DirectRunner' or 'DataFlowRunner'
     project: GCP project name (Required if DataFlowRunner)
     region: GCP region name (Required if DataFlowRunner)
-    tfrutil_path: Path to TFRutil source (Required if DataFlowRunner)
     dataflow_options: Options dict for dataflow runner
     job_label: User supplied description for the beam job name.
     compression: Can be 'gzip' or None for no compression.
@@ -170,7 +166,7 @@ def create_tfrecords(
   df = to_dataframe(input_data, header, names)
 
   _validate_data(df)
-  _validate_runner(df, runner, project, region, tfrutil_path)
+  _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)
@@ -185,7 +181,6 @@ def create_tfrecords(
       runner=runner,
       project=project,
       region=region,
-      tfrutil_path=tfrutil_path,
       output_dir=output_dir,
       compression=compression,
       num_shards=num_shards,

tfrutil/client_test.py

@@ -52,6 +52,8 @@ class InputValidationTest(unittest.TestCase):
 
   def setUp(self):
     self.test_df = test_utils.get_test_df()
+    self.test_region = 'us-central1'
+    self.test_project = 'foo'
 
   def test_valid_dataframe(self):
     """Tests valid DataFrame input."""
@@ -93,29 +95,26 @@ def test_valid_runner(self):
     self.assertIsNone(client._validate_runner(
         self.test_df,
         runner='DirectRunner',
-        project='foo',
-        region='us-central1',
-        tfrutil_path='foo/'))
+        project=self.test_project,
+        region=self.test_region))
 
   def test_invalid_runner(self):
     """Tests invalid runner."""
     with self.assertRaises(AttributeError):
       client._validate_runner(
           self.test_df,
           runner='FooRunner',
-          project='foo',
-          region='us-central1',
-          tfrutil_path='foo/')
+          project=self.test_project,
+          region=self.test_region)
 
   def test_local_path_with_dataflow_runner(self):
     """Tests DataFlowRunner conflict with local path."""
     with self.assertRaises(AttributeError):
       client._validate_runner(
           self.df_test,
           runner='DataFlowRunner',
-          project='foo',
-          region='us-central1',
-          tfrutil_path='foo/')
+          project=self.test_project,
+          region=self.test_region)
 
   def test_gcs_path_with_dataflow_runner(self):
     """Tests DataFlowRunner with gcs path."""
@@ -125,9 +124,8 @@ def test_gcs_path_with_dataflow_runner(self):
         client._validate_runner(
             df2,
             runner='DataFlowRunner',
-            project='foo',
-            region='us-central1',
-            tfrutil_path='foo/'))
+            project=self.test_project,
+            region=self.test_region))
 
   def test_gcs_path_with_dataflow_runner_missing_param(self):
     """Tests DataFlowRunner with missing required parameter."""
@@ -138,8 +136,7 @@ def test_gcs_path_with_dataflow_runner_missing_param(self):
           df2,
           runner='DataFlowRunner',
           project=None,
-          region='us-central1',
-          tfrutil_path='foo/')
+          region=self.test_region)
 
 
 def _make_csv_tempfile(data: List[List[str]]) -> tempfile.NamedTemporaryFile: