Adding TFX IR and node_id as arguments to BaseComponent.

PiperOrigin-RevId: 364924597

Author: tfx-copybara

tfx/dsl/compiler/compiler.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Compiles a TFX pipeline into a TFX DSL IR proto."""
-import json
-import re
 from typing import cast, Iterable, List, Mapping
 
 from tfx import types
@@ -228,12 +226,6 @@ def _compile_node(
           compiler_utils.set_runtime_parameter_pb(
               parameter_value.runtime_parameter, value.name, value.ptype,
               value.default)
-        elif isinstance(value, str) and re.search(
-            data_types.RUNTIME_PARAMETER_PATTERN, value):
-          runtime_param = json.loads(value)
-          compiler_utils.set_runtime_parameter_pb(
-              parameter_value.runtime_parameter, runtime_param.name,
-              runtime_param.ptype, runtime_param.default)
         else:
           try:
             data_types_utils.set_metadata_value(parameter_value.field_value,

tfx/orchestration/kubeflow/base_component.py

@@ -39,6 +39,7 @@
 from tfx.orchestration.kubeflow import utils
 from tfx.orchestration.kubeflow.proto import kubeflow_pb2
 from tfx.orchestration.launcher import base_component_launcher
+from tfx.proto.orchestration import pipeline_pb2
 from tfx.utils import json_utils
 
 from google.protobuf import json_format
@@ -70,6 +71,7 @@ def __init__(
       tfx_image: Text,
       kubeflow_metadata_config: Optional[kubeflow_pb2.KubeflowMetadataConfig],
       component_config: base_component_config.BaseComponentConfig,
+      tfx_ir: pipeline_pb2.Pipeline,
       pod_labels_to_attach: Optional[Dict[Text, Text]] = None):
     """Creates a new Kubeflow-based component.
 
@@ -89,6 +91,7 @@ def __init__(
       kubeflow_metadata_config: Configuration settings for connecting to the
         MLMD store in a Kubeflow cluster.
       component_config: Component config to launch the component.
+      tfx_ir: The TFX intermedia representation of the pipeline.
       pod_labels_to_attach: Optional dict of pod labels to attach to the
         GKE pod.
     """
@@ -117,6 +120,10 @@ def __init__(
         serialized_component,
         '--component_config',
         json_utils.dumps(component_config),
+        '--tfx_ir',
+        json_format.MessageToJson(tfx_ir),
+        '--node_id',
+        component.id,
     ]
 
     if pipeline.enable_cache:

tfx/orchestration/kubeflow/base_component_test.py

@@ -31,6 +31,7 @@
 from tfx.orchestration.kubeflow import base_component
 from tfx.orchestration.kubeflow.proto import kubeflow_pb2
 from tfx.orchestration.launcher import in_process_component_launcher
+from tfx.proto.orchestration import pipeline_pb2
 from tfx.types import channel_utils
 from tfx.types import standard_artifacts
 
@@ -60,6 +61,7 @@ def setUp(self):
 
     self._metadata_config = kubeflow_pb2.KubeflowMetadataConfig()
     self._metadata_config.mysql_db_service_host.environment_variable = 'MYSQL_SERVICE_HOST'
+    self._tfx_ir = pipeline_pb2.Pipeline()
     with dsl.Pipeline('test_pipeline'):
       self.component = base_component.BaseComponent(
           component=statistics_gen,
@@ -72,6 +74,7 @@ def setUp(self):
           tfx_image='container_image',
           kubeflow_metadata_config=self._metadata_config,
           component_config=None,
+          tfx_ir=self._tfx_ir,
       )
     self.tfx_component = statistics_gen
 
@@ -159,6 +162,7 @@ def setUp(self):
 
     self._metadata_config = kubeflow_pb2.KubeflowMetadataConfig()
     self._metadata_config.mysql_db_service_host.environment_variable = 'MYSQL_SERVICE_HOST'
+    self._tfx_ir = pipeline_pb2.Pipeline()
     with dsl.Pipeline('test_pipeline'):
       self.example_gen = base_component.BaseComponent(
           component=example_gen,
@@ -170,7 +174,8 @@ def setUp(self):
           pipeline_root=test_pipeline_root,
           tfx_image='container_image',
           kubeflow_metadata_config=self._metadata_config,
-          component_config=None)
+          component_config=None,
+          tfx_ir=self._tfx_ir)
       self.statistics_gen = base_component.BaseComponent(
           component=statistics_gen,
           component_launcher_class=in_process_component_launcher
@@ -182,6 +187,7 @@ def setUp(self):
           tfx_image='container_image',
           kubeflow_metadata_config=self._metadata_config,
           component_config=None,
+          tfx_ir=self._tfx_ir
       )
 
     self.tfx_example_gen = example_gen
@@ -221,6 +227,10 @@ def testContainerOpArguments(self):
         formatted_statistics_gen,
         '--component_config',
         'null',
+        '--tfx_ir',
+        '{}',
+        '--node_id',
+        'StatisticsGen.foo',
     ]
     example_gen_expected_args = [
         '--pipeline_name',
@@ -243,6 +253,10 @@ def testContainerOpArguments(self):
         formatted_example_gen,
         '--component_config',
         'null',
+        '--tfx_ir',
+        '{}',
+        '--node_id',
+        'CsvExampleGen',
     ]
     try:
       self.assertEqual(

tfx/orchestration/kubeflow/kubeflow_dag_runner.py

@@ -19,14 +19,15 @@
 
 import os
 import re
-from typing import Callable, Dict, List, Optional, Text, Type
+from typing import Callable, Dict, List, Optional, Text, Type, cast
 
+from absl import logging
 from kfp import compiler
 from kfp import dsl
 from kfp import gcp
 from kubernetes import client as k8s_client
-
 from tfx import version
+from tfx.dsl.compiler import compiler as tfx_compiler
 from tfx.orchestration import data_types
 from tfx.orchestration import pipeline as tfx_pipeline
 from tfx.orchestration import tfx_runner
@@ -38,9 +39,11 @@
 from tfx.orchestration.launcher import base_component_launcher
 from tfx.orchestration.launcher import in_process_component_launcher
 from tfx.orchestration.launcher import kubernetes_component_launcher
+from tfx.proto.orchestration import pipeline_pb2
 from tfx.utils import json_utils
 from tfx.utils import telemetry_utils
 
+
 # OpFunc represents the type of a function that takes as input a
 # dsl.ContainerOp and returns the same object. Common operations such as adding
 # k8s secrets, mounting volumes, specifying the use of TPUs and so on can be
@@ -249,9 +252,11 @@ def __init__(
     if config and not isinstance(config, KubeflowDagRunnerConfig):
       raise TypeError('config must be type of KubeflowDagRunnerConfig.')
     super(KubeflowDagRunner, self).__init__(config or KubeflowDagRunnerConfig())
+    self._config = cast(KubeflowDagRunnerConfig, self._config)
     self._output_dir = output_dir or os.getcwd()
     self._output_filename = output_filename
     self._compiler = compiler.Compiler()
+    self._tfx_compiler = tfx_compiler.Compiler()
     self._params = []  # List of dsl.PipelineParam used in this pipeline.
     self._deduped_parameter_names = set()  # Set of unique param names used.
     if pod_labels_to_attach is None:
@@ -307,6 +312,7 @@ def _construct_pipeline_graph(self, pipeline: tfx_pipeline.Pipeline,
       pipeline_root: dsl.PipelineParam representing the pipeline root.
     """
     component_to_kfp_op = {}
+    tfx_ir = self._generate_tfx_ir(pipeline)
 
     # Assumption: There is a partial ordering of components in the list, i.e.,
     # if component A depends on component B and C, then A appears after B and C
@@ -332,13 +338,20 @@ def _construct_pipeline_graph(self, pipeline: tfx_pipeline.Pipeline,
           tfx_image=self._config.tfx_image,
           kubeflow_metadata_config=self._config.kubeflow_metadata_config,
           component_config=component_config,
-          pod_labels_to_attach=self._pod_labels_to_attach)
+          pod_labels_to_attach=self._pod_labels_to_attach,
+          tfx_ir=tfx_ir)
 
       for operator in self._config.pipeline_operator_funcs:
         kfp_component.container_op.apply(operator)
 
       component_to_kfp_op[component] = kfp_component.container_op
 
+  def _generate_tfx_ir(
+      self, pipeline: tfx_pipeline.Pipeline) -> pipeline_pb2.Pipeline:
+    result = self._tfx_compiler.compile(pipeline)
+    logging.info('Generated pipeline:\n %s', result)
+    return result
+
   def run(self, pipeline: tfx_pipeline.Pipeline):
     """Compiles and outputs a Kubeflow Pipeline YAML definition file.