code: def test_subpipeline_task_scheduler(self, cancel_pipeline):
      sleep_time = (
          subpipeline_task_scheduler._POLLING_INTERVAL_SECS.value * 5
      )
  
      with self._mlmd_connection as mlmd_connection:
          test_utils.fake_example_gen_run(
              mlmd_connection, self._example_gen, 1, 1
          )
  
          [sub_pipeline_task] = test_utils.run_generator_and_test(
              test_case=self,
              mlmd_connection_manager=self._mlmd_cm,
              generator_class=sptg.SyncPipelineTaskGenerator,
              pipeline=self._pipeline,
              task_queue=self._task_queue,
              use_task_queue=True,
              service_job_manager=None,
              num_initial_executions=1,
              num_tasks_generated=1,
              num_new_executions=1,
              num_active_executions=1,
              expected_exec_nodes=[self._sub_pipeline],
              ignore_update_node_state_tasks=True,
              expected_context_names=[
                  "my_sub_pipeline",
                  f"my_sub_pipeline_{self._pipeline_run_id}",
                  "my_pipeline",
                  self._pipeline_run_id,
                  "my_sub_pipeline.my_sub_pipeline",
              ],
          )
  
          # There should be only 1 orchestrator execution for the outer pipeline.
          pipeline_states = pstate.PipelineState.load_all_active(
              mlmd_connection
          )
          self.assertLen(pipeline_states, 1)
  
          ts_result = []
          scheduler = (
              subpipeline_task_scheduler.SubPipelineTaskScheduler(
                  mlmd_handle=mlmd_connection,
                  pipeline=self._pipeline,
                  task=sub_pipeline_task,
              )
          )
  
          def start_scheduler(ts_result):
              ts_result.append(scheduler.schedule())
  
          threading.Thread(
              target=start_scheduler, args=(ts_result,)
          ).start()
  
          # Wait for sometime for the update to go through.
          time.sleep(sleep_time)
  
          # There should be another orchestrator execution for the inner pipeline.
          pipeline_states = pstate.PipelineState.load_all_active(
              mlmd_connection
          )
          self.assertLen(pipeline_states, 2)
          subpipeline_state = pstate.PipelineState.load(
              mlmd_connection,
              task_lib.PipelineUid(pipeline_id="my_sub_pipeline"),
          )
  
          # The scheduler is still waiting for subpipeline to finish.
          self.assertEmpty(ts_result)
  
          if cancel_pipeline:
              # Call cancel() to initiate the cancel.
              scheduler.cancel(
                  task_lib.CancelNodeTask(
                      node_uid=task_lib.NodeUid.from_node(
                          self._pipeline,
                          self._sub_pipeline,
                      )
                  )
              )
  
              # Sets the cancel state on subpipeline.
              def _cancel(pipeline_state):
                  time.sleep(2.0)
                  with pipeline_state:
                      if pipeline_state.is_stop_initiated():
                          pipeline_state.set_pipeline_execution_state(
                              metadata_store_pb2.Execution.CANCELED
                          )
  
              threading.Thread(
                  target=_cancel, args=(subpipeline_state,)
              ).start()
  
              # Wait for the update to go through.
              time.sleep(sleep_time)
  
              self.assertLen(ts_result, 1)
              self.assertEqual(
                  status_lib.Code.CANCELLED, ts_result[0].status.code
              )
          else:
              # Mark inner pipeline as COMPLETE.
              def _complete(pipeline_state):
                  with pipeline_state:
                      pipeline_state.set_pipeline_execution_state(
                          metadata_store_pb2.Execution.COMPLETE
                      )
  
              threading.Thread(
                  target=_complete, args=(subpipeline_state,)
              ).start()
  
              # Wait for the update to go through.
              time.sleep(sleep_time)
  
              self.assertLen(ts_result, 1)
              self.assertEqual(
                  status_lib.Code.OK, ts_result[0].status.code
              )
              self.assertIsInstance(
                  ts_result[0].output, ts.ExecutorNodeOutput
              )
  
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 4 for calling the SubpipelineTaskSchedulerTest obj's underlying member method code with example usage: obj.test_subpipeline_task_scheduler(cancel_pipeline) without return types