Minor fixes to unit tests for orchestration/portable

PiperOrigin-RevId: 327320340

Author: tfx-copybara

RELEASE.md

@@ -66,6 +66,7 @@
     python (2.7/3.5/3.6) is not available anymore in `tensorflow/tfx` images
     on docker hub. Virtualenv is not used anymore.
 *   Depends on `pyarrow>=0.17,<0.18`.
+*   Depends on `attrs>=19.3.0,<20`.
 
 ## Breaking changes
 *   Changed the URIs of the value artifacts to point to files.

tfx/dependencies.py

@@ -24,6 +24,7 @@ def make_required_install_packages():
       # LINT.IfChange
       'apache-beam[gcp]>=2.22,<3',
       # LINT.ThenChange(examples/chicago_taxi_pipeline/setup/setup_beam.sh)
+      'attrs>=19.3.0,<20',
       'click>=7,<8',
       'docker>=4.1,<5',
       'google-api-python-client>=1.7.8,<2',

tfx/orchestration/portable/base_executor_operator.py

@@ -14,7 +14,7 @@
 """Base class to define how to operator an executor."""
 
 import abc
-from typing import Any, Dict, List, Optional, Text
+from typing import Any, Dict, List, Optional
 
 import attr
 import six
@@ -30,30 +30,31 @@
 # TODO(b/150979622): We should introduce an id that is not changed across
 # retires of the same component run and pass it to executor operators for
 # human-readability purpose.
-@attr.s(auto_attribs=True)
+# TODO(b/165359991): Restore 'auto_attribs=True' once we drop Python3.5 support.
+@attr.s
 class ExecutionInfo:
   """A struct to store information for an execution."""
   # The metadata of this execution that is registered in MLMD.
-  execution_metadata: metadata_store_pb2.Execution = None
+  execution_metadata = attr.ib(type=metadata_store_pb2.Execution, default=None)
   # The input map to feed to executor
-  input_dict: Dict[Text, List[types.Artifact]] = None
+  input_dict = attr.ib(type=Dict[str, List[types.Artifact]], default=None)
   # The output map to feed to executor
-  output_dict: Dict[Text, List[types.Artifact]] = None
+  output_dict = attr.ib(type=Dict[str, List[types.Artifact]], default=None)
   # The exec_properties to feed to executor
-  exec_properties: Dict[Text, Any] = None
+  exec_properties = attr.ib(type=Dict[str, Any], default=None)
   # The uri to executor result, note that Executors and Launchers may not run
   # in the same process, so executors should use this uri to "return"
   # ExecutorOutput to the launcher.
-  executor_output_uri: Text = None
+  executor_output_uri = attr.ib(type=str, default=None)
   # Stateful working dir will be deterministic given pipeline, node and run_id.
   # The typical usecase is to restore long running executor's state after
   # eviction. For examples, a Trainer can use this directory to store
   # checkpoints.
-  stateful_working_dir: Text = None
+  stateful_working_dir = attr.ib(type=str, default=None)
   # The config of this Node.
-  pipeline_node: pipeline_pb2.PipelineNode = None
+  pipeline_node = attr.ib(type=pipeline_pb2.PipelineNode, default=None)
   # The config of the pipeline that this node is running in.
-  pipeline_info: pipeline_pb2.PipelineInfo = None
+  pipeline_info = attr.ib(type=pipeline_pb2.PipelineInfo, default=None)
 
 
 class BaseExecutorOperator(six.with_metaclass(abc.ABCMeta, object)):

tfx/orchestration/portable/beam_dag_runner_test.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for tfx.orchestration.portable.beam_dag_runner."""
-
+import os
 import mock
 import tensorflow as tf
+
 from tfx.orchestration import metadata
 from tfx.orchestration.portable import beam_dag_runner
 from tfx.orchestration.portable import test_utils
@@ -46,8 +47,10 @@ def setUp(self):
     super(BeamDagRunnerTest, self).setUp()
     # Setup pipelines
     self._pipeline = pipeline_pb2.Pipeline()
-    self.load_proto_from_text('pipeline_for_launcher_test.pbtxt',
-                              self._pipeline)
+    self.load_proto_from_text(
+        os.path.join(
+            os.path.dirname(__file__), 'testdata',
+            'pipeline_for_launcher_test.pbtxt'), self._pipeline)
 
   @mock.patch.multiple(
       beam_dag_runner,

tfx/orchestration/portable/inputs_utils_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for tfx.orchestration.portable.inputs_utils."""
+import os
 import tensorflow as tf
 
 from tfx import types
@@ -28,6 +29,10 @@
 
 class InputsUtilsTest(test_utils.TfxTest):
 
+  def setUp(self):
+    super().setUp()
+    self._testdata_dir = os.path.join(os.path.dirname(__file__), 'testdata')
+
   def testResolveParameters(self):
     parameters = pipeline_pb2.NodeParameters()
     text_format.Parse(
@@ -65,8 +70,9 @@ def testResolveParametersFail(self):
 
   def testResolverInputsArtifacts(self):
     pipeline = pipeline_pb2.Pipeline()
-    self.load_proto_from_text('pipeline_for_input_resolver_test.pbtxt',
-                              pipeline)
+    self.load_proto_from_text(
+        os.path.join(self._testdata_dir,
+                     'pipeline_for_input_resolver_test.pbtxt'), pipeline)
     my_example_gen = pipeline.nodes[0].pipeline_node
     another_example_gen = pipeline.nodes[1].pipeline_node
     my_transform = pipeline.nodes[2].pipeline_node

tfx/orchestration/portable/launcher.py

@@ -45,21 +45,23 @@
 }
 
 
-@attr.s(auto_attribs=True)
+# TODO(b/165359991): Restore 'auto_attribs=True' once we drop Python3.5 support.
+@attr.s
 class _PrepareExecutionResult:
   """A wrapper class using as the return value of _prepare_execution()."""
 
   # The information used by executor operators.
-  execution_info: Optional[base_executor_operator.ExecutionInfo]
+  execution_info = attr.ib(
+      type=base_executor_operator.ExecutionInfo, default=None)
   # Contexts of the execution, usually used by Publisher.
-  contexts: List[metadata_store_pb2.Context]
+  contexts = attr.ib(type=List[metadata_store_pb2.Context], default=None)
   # TODO(b/156126088): Update the following documentation when this bug is
   # closed.
   # Whether an execution is needed. An execution is not needed when:
   # 1) Not all the required input are ready.
   # 2) The input value doesn't meet the driver's requirement.
   # 3) Cache result is used.
-  is_execution_needed: bool = False
+  is_execution_needed = attr.ib(type=bool, default=False)
 
 
 class Launcher(object):
@@ -241,7 +243,7 @@ def _publish_failed_execution(
           metadata_handler=m, execution_id=execution_id, contexts=contexts)
 
   def _clean_up(self, execution_info: base_executor_operator.ExecutionInfo):
-    tf.io.gfile.remove(execution_info.stateful_working_dir)
+    tf.io.gfile.rmtree(execution_info.stateful_working_dir)
 
   def launch(self) -> Optional[metadata_store_pb2.Execution]:
     """Executes the component, includes driver, executor and publisher.

tfx/orchestration/portable/launcher_test.py

@@ -67,10 +67,14 @@ def setUp(self):
     connection_config.sqlite.SetInParent()
     self._mlmd_connection = metadata.Metadata(
         connection_config=connection_config)
+    self._testdata_dir = os.path.join(os.path.dirname(__file__), 'testdata')
 
     # Sets up pipelines
     pipeline = pipeline_pb2.Pipeline()
-    self.load_proto_from_text('pipeline_for_launcher_test.pbtxt', pipeline)
+    self.load_proto_from_text(
+        os.path.join(
+            os.path.dirname(__file__), 'testdata',
+            'pipeline_for_launcher_test.pbtxt'), pipeline)
     self._pipeline_info = pipeline.pipeline_info
     self._pipeline_runtime_spec = pipeline.runtime_spec
     self._pipeline_runtime_spec.pipeline_root.field_value.string_value = (

tfx/orchestration/portable/mlmd/context_lib_test.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for tfx.orchestration.portable.mlmd.context_lib."""
+import os
 import tensorflow as tf
+
 from tfx.orchestration import metadata
 from tfx.orchestration.portable import test_utils
 from tfx.orchestration.portable.mlmd import context_lib
@@ -26,10 +28,14 @@ def setUp(self):
     super().setUp()
     self._connection_config = metadata_store_pb2.ConnectionConfig()
     self._connection_config.sqlite.SetInParent()
+    self._testdata_dir = os.path.join(
+        os.path.dirname(os.path.dirname(__file__)), 'testdata')
 
   def testRegisterContexts(self):
     node_contexts = pipeline_pb2.NodeContexts()
-    self.load_proto_from_text('node_context_spec.pbtxt', node_contexts)
+    self.load_proto_from_text(
+        os.path.join(self._testdata_dir, 'node_context_spec.pbtxt'),
+        node_contexts)
     with metadata.Metadata(connection_config=self._connection_config) as m:
       context_lib.register_contexts_if_not_exists(
           metadata_handler=m, node_contexts=node_contexts)

tfx/orchestration/portable/mlmd/execution_lib_test.py

@@ -96,35 +96,17 @@ def testPrepareExecution(self):
           """, result)
 
   def testArtifactAndEventPairs(self):
-    model = standard_artifacts.Model()
-    model.uri = 'model'
     example = standard_artifacts.Examples()
     example.uri = 'example'
     example.id = 1
 
-    expected_artifact_one = metadata_store_pb2.Artifact()
-    expected_artifact_two = metadata_store_pb2.Artifact()
-    text_format.Parse("""
-        type_id: 1
-        uri: 'model'""", expected_artifact_one)
+    expected_artifact = metadata_store_pb2.Artifact()
     text_format.Parse(
         """
         id: 1
-        type_id: 2
-        uri: 'example'""", expected_artifact_two)
-    expected_event_one = metadata_store_pb2.Event()
-    expected_event_two = metadata_store_pb2.Event()
-    text_format.Parse(
-        """
-        path {
-          steps {
-            key: 'model'
-          }
-          steps {
-            index: 0
-          }
-        }
-        type: INPUT""", expected_event_one)
+        type_id: 1
+        uri: 'example'""", expected_artifact)
+    expected_event = metadata_store_pb2.Event()
     text_format.Parse(
         """
         path {
@@ -135,18 +117,15 @@ def testArtifactAndEventPairs(self):
             index: 0
           }
         }
-        type: INPUT""", expected_event_two)
+        type: INPUT""", expected_event)
 
     with metadata.Metadata(connection_config=self._connection_config) as m:
       result = execution_lib._create_artifact_and_event_pairs(
           m, {
-              'model': [model],
               'example': [example],
           }, metadata_store_pb2.Event.INPUT)
 
-      self.assertListEqual([(expected_artifact_one, expected_event_one),
-                            (expected_artifact_two, expected_event_two)],
-                           result)
+      self.assertCountEqual([(expected_artifact, expected_event)], result)
 
   def testPutExecutionGraph(self):
     with metadata.Metadata(connection_config=self._connection_config) as m:

tfx/orchestration/portable/outputs_utils.py

@@ -38,7 +38,10 @@ def remove_output_dirs(output_dict: Dict[Text, List[types.Artifact]]) -> None:
   """Remove dirs of output artifacts' URI."""
   for _, artifact_list in output_dict.items():
     for artifact in artifact_list:
-      tf.io.gfile.remove(artifact.uri)
+      if tf.io.gfile.isdir(artifact.uri):
+        tf.io.gfile.rmtree(artifact.uri)
+      else:
+        tf.io.gfile.remove(artifact.uri)
 
 
 class OutputsResolver:

tfx/orchestration/portable/runtime_parameter_utils_test.py

@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for tfx.orchestration.portable.runtime_parameter_utils."""
+import os
+
 import tensorflow as tf
 
-from ml_metadata.proto import metadata_store_pb2
 from tfx.orchestration.portable import runtime_parameter_utils
 from tfx.orchestration.portable import test_utils
 from tfx.proto.orchestration import pipeline_pb2
+from ml_metadata.proto import metadata_store_pb2
 
 
 class RuntimeParameterUtilsTest(test_utils.TfxTest):
@@ -26,13 +28,18 @@ def setUp(self):
     super().setUp()
     self._connection_config = metadata_store_pb2.ConnectionConfig()
     self._connection_config.sqlite.SetInParent()
+    self._testdata_dir = os.path.join(os.path.dirname(__file__), 'testdata')
 
   def testFullySubstituteRuntimeParameter(self):
     pipeline = pipeline_pb2.Pipeline()
     expected = pipeline_pb2.Pipeline()
-    self.load_proto_from_text('pipeline_with_runtime_parameter.pbtxt', pipeline)
     self.load_proto_from_text(
-        'pipeline_with_runtime_parameter_substituted.pbtxt', expected)
+        os.path.join(self._testdata_dir,
+                     'pipeline_with_runtime_parameter.pbtxt'), pipeline)
+    self.load_proto_from_text(
+        os.path.join(self._testdata_dir,
+                     'pipeline_with_runtime_parameter_substituted.pbtxt'),
+        expected)
     runtime_parameter_utils.substitute_runtime_parameter(
         pipeline, {
             'context_name_rp': 'my_context',
@@ -44,9 +51,14 @@ def testFullySubstituteRuntimeParameter(self):
   def testPartiallySubstituteRuntimeParameter(self):
     pipeline = pipeline_pb2.Pipeline()
     expected = pipeline_pb2.Pipeline()
-    self.load_proto_from_text('pipeline_with_runtime_parameter.pbtxt', pipeline)
     self.load_proto_from_text(
-        'pipeline_with_runtime_parameter_partially_substituted.pbtxt', expected)
+        os.path.join(self._testdata_dir,
+                     'pipeline_with_runtime_parameter.pbtxt'), pipeline)
+    self.load_proto_from_text(
+        os.path.join(
+            self._testdata_dir,
+            'pipeline_with_runtime_parameter_partially_substituted.pbtxt'),
+        expected)
     runtime_parameter_utils.substitute_runtime_parameter(
         pipeline, {
             'context_name_rp': 'my_context',
@@ -55,7 +67,9 @@ def testPartiallySubstituteRuntimeParameter(self):
 
   def testSubstituteRuntimeParameterFail(self):
     pipeline = pipeline_pb2.Pipeline()
-    self.load_proto_from_text('pipeline_with_runtime_parameter.pbtxt', pipeline)
+    self.load_proto_from_text(
+        os.path.join(self._testdata_dir,
+                     'pipeline_with_runtime_parameter.pbtxt'), pipeline)
     with self.assertRaisesRegex(RuntimeError, 'Runtime parameter type'):
       runtime_parameter_utils.substitute_runtime_parameter(
           pipeline,

tfx/orchestration/portable/test_utils.py

@@ -34,10 +34,9 @@ def setUp(self):
         self._testMethodName)
     tf.io.gfile.makedirs(self.tmp_dir)
 
-  def load_proto_from_text(self, file_name: Text,
+  def load_proto_from_text(self, path: Text,
                            proto_message: message.Message) -> message.Message:
     """Loads proto message from serialized text."""
-    path = os.path.join(os.path.dirname(__file__), 'testdata', file_name)
     return io_utils.parse_pbtxt_file(path, proto_message)
 
   def assertProtoPartiallyEquals(