code: def test_subpipeline_ir_rewrite(self):
      old_ir = copy.deepcopy(self._sub_pipeline.raw_proto())
      new_ir = subpipeline_task_scheduler.subpipeline_ir_rewrite(
          self._sub_pipeline.raw_proto(), execution_id=42
      )
  
      # Asserts original IR is unmodified.
      self.assertProtoEquals(self._sub_pipeline.raw_proto(), old_ir)
  
      # Asserts begin node has no upstream and end node has no downstream.
      self.assertEmpty(new_ir.nodes[0].pipeline_node.upstream_nodes)
      self.assertEmpty(new_ir.nodes[-1].pipeline_node.downstream_nodes)
  
      # New run id should be <old_run_id>_<execution_id>.
      old_run_id = (
          old_ir.runtime_spec.pipeline_run_id.field_value.string_value
      )
      new_run_id = (
          new_ir.runtime_spec.pipeline_run_id.field_value.string_value
      )
      self.assertEqual(new_run_id, old_run_id + "_42")
  
      # All nodes should associate with the new pipeline run id.
      for node in new_ir.nodes:
          pipeline_run_context_names = set()
          for c in node.pipeline_node.contexts.contexts:
              if c.type.name == "pipeline_run":
                  pipeline_run_context_names.add(
                      c.name.field_value.string_value
                  )
          self.assertIn(new_run_id, pipeline_run_context_names)
          self.assertNotIn(old_run_id, pipeline_run_context_names)
  
      # All inputs except those of PipelineBeginNode's should associate with the
      # new pipeline run id.
      for node in new_ir.nodes[1:]:
          for input_spec in node.pipeline_node.inputs.inputs.values():
              for channel in input_spec.channels:
                  pipeline_run_context_names = set()
                  for context_query in channel.context_queries:
                      if context_query.type.name == "pipeline_run":
                          pipeline_run_context_names.add(
                              context_query.name.field_value.string_value
                          )
                  self.assertIn(new_run_id, pipeline_run_context_names)
                  self.assertNotIn(
                      old_run_id, pipeline_run_context_names
                  )
  
imports: import copy
  import os
  import threading
  import time
  import uuid
  from absl.testing import flagsaver
  from absl.testing import parameterized
  import tensorflow
  from tfx.dsl.compiler import constants
  from tfx.orchestration.experimental.core import pipeline_state
  from tfx.orchestration.experimental.core import sync_pipeline_task_gen
  from tfx.orchestration.experimental.core import task
  from tfx.orchestration.experimental.core import task_queue
  from tfx.orchestration.experimental.core import task_scheduler
  from tfx.orchestration.experimental.core import test_utils
  from tfx.orchestration.experimental.core.task_schedulers import subpipeline_task_scheduler
  from tfx.orchestration.experimental.core.testing import test_subpipeline
  from tfx.orchestration import mlmd_connection_manager
  from tfx.orchestration.portable import runtime_parameter_utils
  from tfx.utils import status
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 3 for calling the SubpipelineTaskSchedulerTest obj's underlying member method code with example usage: obj.test_subpipeline_ir_rewrite() without return types