Update inferred setup.py path to full path.

Change-Id: I50522de73ed523f248b3894022ce3bba7aac611e

Author: cfezequiel

tfrutil/beam_pipeline.py

@@ -34,6 +34,16 @@
 from tfrutil import constants
 
 
+def _get_setup_py_filepath() -> str:
+  """Returns the file path to the setup.py file.
+
+  The location of the setup.py file is needed to run Dataflow jobs.
+  """
+
+  return os.path.join(
+      os.path.dirname(os.path.abspath(__file__)), '..', 'setup.py')
+
+
 def _get_job_name(job_label: str = None) -> str:
   """Returns Beam runner job name.
 
@@ -84,7 +94,7 @@ def _get_pipeline_options(
   if region:
     options_dict['region'] = region
   if runner == 'DataFlowRunner':
-    options_dict['setup_file'] = os.path.join('..', 'setup.py')
+    options_dict['setup_file'] = _get_setup_py_filepath()
   if dataflow_options:
     options_dict.update(dataflow_options)
 

tfrutil/beam_pipeline_test.py

@@ -16,6 +16,7 @@
 
 """Tests for beam_pipeline."""
 
+import os
 import unittest
 from unittest import mock
 
@@ -28,7 +29,7 @@
 # pylint: disable=protected-access
 
 class BeamPipelineTests(unittest.TestCase):
-  """Tests for beam_image.py"""
+  """Tests for beam_pipeline.py"""
 
   def test_processing_fn_with_int_label(self):
     'Test preprocessing fn with integer label.'
@@ -76,3 +77,9 @@ def test_partition_fn(self):
       self.assertEqual(
           index, i,
           '{} should be index {} but was index {}'.format(part, i, index))
+
+  def test_get_setup_py_filepath(self):
+    """Tests `_get_setup_py_filepath`."""
+    filepath = beam_pipeline._get_setup_py_filepath()
+    self.assertTrue(os.path.isfile(filepath))
+    self.assertTrue(os.path.isabs(filepath))