Encode producer component id and output key when CWP is created from …

…an OutputChannel PiperOrigin-RevId: 619667393
tensorflow · Mar 28, 2024 · 5ef7ee6 · 5ef7ee6
1 parent 92fef51
commit 5ef7ee6
Show file tree

Hide file tree

Showing 9 changed files with 41 additions and 23 deletions.
diff --git a/tfx/dsl/compiler/compiler_test.py b/tfx/dsl/compiler/compiler_test.py
@@ -205,10 +205,17 @@ def testCompileAdditionalCustomPropertyNameConflictError(self):
   def testCompileDynamicExecPropTypeError(self):
     dsl_compiler = compiler.Compiler()
     test_pipeline = dynamic_exec_properties_pipeline.create_test_pipeline()
+    upstream_component = next(
+        c
+        for c in test_pipeline.components
+        if isinstance(c, dynamic_exec_properties_pipeline.UpstreamComponent)
+    )
     downstream_component = next(
         c for c in test_pipeline.components
         if isinstance(c, dynamic_exec_properties_pipeline.DownstreamComponent))
-    test_wrong_type_channel = channel.Channel(_MyType).future().value
+    test_wrong_type_channel = (
+        channel.OutputChannel(_MyType, upstream_component, "foo").future().value
+    )
     downstream_component.exec_properties["input_num"] = test_wrong_type_channel
     with self.assertRaisesRegex(
         ValueError, ".*channel must be of a value artifact type.*"

diff --git a/tfx/dsl/compiler/testdata/composable_pipeline_async_input_v2_ir.pbtxt b/tfx/dsl/compiler/testdata/composable_pipeline_async_input_v2_ir.pbtxt
@@ -1886,7 +1886,7 @@ nodes {
                               index_op {
                                 expression {
                                   placeholder {
-                                    key: "blessing"
+                                    key: "Evaluator_blessing"
                                   }
                                 }
                               }
@@ -2983,7 +2983,7 @@ nodes {
                                     index_op {
                                       expression {
                                         placeholder {
-                                          key: "_infra-validator-pipeline.blessing"
+                                          key: "infra-validator-pipeline_blessing"
                                         }
                                       }
                                     }

diff --git a/tfx/dsl/compiler/testdata/composable_pipeline_input_v2_ir.pbtxt b/tfx/dsl/compiler/testdata/composable_pipeline_input_v2_ir.pbtxt
@@ -82,7 +82,8 @@ nodes {
           }
         }
         execution_options {
-          caching_options {}
+          caching_options {
+          }
         }
       }
     }
@@ -1169,7 +1170,8 @@ nodes {
         upstream_nodes: "data-ingestion-pipeline"
         downstream_nodes: "Trainer"
         execution_options {
-          caching_options {}
+          caching_options {
+          }
           strategy: LAZILY_ALL_UPSTREAM_NODES_SUCCEEDED
           max_execution_retries: 10
         }
@@ -2143,7 +2145,7 @@ nodes {
                               index_op {
                                 expression {
                                   placeholder {
-                                    key: "_Evaluator.blessing"
+                                    key: "Evaluator_blessing"
                                   }
                                 }
                               }
@@ -2206,7 +2208,8 @@ nodes {
         downstream_nodes: "Pusher"
         downstream_nodes: "infra-validator-pipeline"
         execution_options {
-          caching_options {}
+          caching_options {
+          }
         }
       }
     }
@@ -2507,7 +2510,8 @@ nodes {
             upstream_nodes: "validate-and-push-pipeline_begin"
             downstream_nodes: "InfraValidator"
             execution_options {
-              caching_options {}
+              caching_options {
+              }
             }
           }
         }
@@ -3347,7 +3351,7 @@ nodes {
                                     index_op {
                                       expression {
                                         placeholder {
-                                          key: "_infra-validator-pipeline.blessing"
+                                          key: "infra-validator-pipeline_blessing"
                                         }
                                       }
                                     }

diff --git a/tfx/dsl/compiler/testdata/conditional_pipeline_input_v2_ir.pbtxt b/tfx/dsl/compiler/testdata/conditional_pipeline_input_v2_ir.pbtxt
@@ -1001,7 +1001,7 @@ nodes {
                           index_op {
                             expression {
                               placeholder {
-                                key: "_Evaluator.blessing"
+                                key: "Evaluator_blessing"
                               }
                             }
                           }
@@ -1264,7 +1264,7 @@ nodes {
                           index_op {
                             expression {
                               placeholder {
-                                key: "_Evaluator.blessing"
+                                key: "Evaluator_blessing"
                               }
                             }
                           }
@@ -1301,7 +1301,7 @@ nodes {
                                 index_op {
                                   expression {
                                     placeholder {
-                                      key: "_InfraValidator.blessing"
+                                      key: "InfraValidator_blessing"
                                     }
                                   }
                                 }
@@ -1333,7 +1333,7 @@ nodes {
                                       index_op {
                                         expression {
                                           placeholder {
-                                            key: "model"
+                                            key: "Trainer_model"
                                           }
                                         }
                                       }

diff --git a/tfx/dsl/compiler/testdata/dynamic_exec_properties_pipeline_input_v2_ir.pbtxt b/tfx/dsl/compiler/testdata/dynamic_exec_properties_pipeline_input_v2_ir.pbtxt
@@ -180,7 +180,7 @@ nodes {
                     index_op {
                       expression {
                         placeholder {
-                          key: "_UpstreamComponent.num"
+                          key: "UpstreamComponent_num"
                         }
                       }
                     }

diff --git a/tfx/dsl/compiler/testdata/pipeline_with_annotations_input_v2_ir.pbtxt b/tfx/dsl/compiler/testdata/pipeline_with_annotations_input_v2_ir.pbtxt
@@ -221,7 +221,7 @@ nodes {
                     index_op {
                       expression {
                         placeholder {
-                          key: "_UpstreamComponent.num"
+                          key: "UpstreamComponent_num"
                         }
                       }
                     }

diff --git a/tfx/types/channel.py b/tfx/types/channel.py
@@ -204,7 +204,7 @@ def trigger_by_property(self, *property_keys: str):
     return self._with_input_trigger(TriggerByProperty(property_keys))
 
   def future(self) -> ChannelWrappedPlaceholder:
-    return ChannelWrappedPlaceholder(self)
+    raise NotImplementedError()
 
   def __eq__(self, other):
     return self is other
@@ -557,6 +557,11 @@ def set_external(self, predefined_artifact_uris: List[str]) -> None:
   def set_as_async_channel(self) -> None:
     self._is_async = True
 
+  def future(self) -> ChannelWrappedPlaceholder:
+    return ChannelWrappedPlaceholder(
+        self, f'{self.producer_component_id}_{self.output_key}'
+    )
+
 
 @doc_controls.do_not_generate_docs
 class UnionChannel(BaseChannel):
@@ -703,6 +708,9 @@ def trigger_by_property(self, *property_keys: str):
         'trigger_by_property is not implemented for PipelineInputChannel.'
     )
 
+  def future(self) -> ChannelWrappedPlaceholder:
+    return ChannelWrappedPlaceholder(self, f'{self._output_key}')
+
 
 class ExternalPipelineChannel(BaseChannel):
   """Channel subtype that is used to get artifacts from external MLMD db."""
@@ -787,7 +795,8 @@ def set_key(self, key: Optional[str]):
     Args:
       key: The new key for the channel.
     """
-    self._key = key
+    del key  # unused.
+    return
 
   def __getitem__(self, index: int) -> ChannelWrappedPlaceholder:
     if self._index is not None:

diff --git a/tfx/types/channel_utils.py b/tfx/types/channel_utils.py
@@ -211,10 +211,8 @@ def unwrap_simple_channel_placeholder(
       # proto paths above and been getting default messages all along. If this
       # sub-message is present, then the whole chain was correct.
       not index_op.expression.HasField('placeholder')
-      # ChannelWrappedPlaceholder uses INPUT_ARTIFACT for some reason, and has
-      # no key when encoded with encode().
+      # ChannelWrappedPlaceholder uses INPUT_ARTIFACT for some reason.
       or cwp.type != placeholder_pb2.Placeholder.Type.INPUT_ARTIFACT
-      or cwp.key
       # For the `[0]` part of the desired shape.
       or index_op.index != 0
   ):

diff --git a/tfx/types/channel_utils_test.py b/tfx/types/channel_utils_test.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Tests for tfx.utils.channel."""
 
-import tensorflow as tf
+from absl.testing import absltest
 from tfx.dsl.placeholder import placeholder as ph
 from tfx.types import artifact
 from tfx.types import channel
@@ -25,7 +25,7 @@ class _MyArtifact(artifact.Artifact):
   TYPE_NAME = 'MyTypeName'
 
 
-class ChannelUtilsTest(tf.test.TestCase):
+class ChannelUtilsTest(absltest.TestCase):
 
   def testArtifactCollectionAsChannel(self):
     instance_a = _MyArtifact()
@@ -125,4 +125,4 @@ def testUnwrapSimpleChannelPlaceholderRejectsComplexPlaceholders(self):
 
 
 if __name__ == '__main__':
-  tf.test.main()
+  absltest.main()