@@ -180,9 +180,11 @@ def teardown_training_pipeline(shared_state, pipeline_client):
 )
 
 # Waiting for training pipeline to be in CANCELLED state
+timeout = shared_state["cancel_batch_prediction_job_timeout"]
 helpers.wait_for_job_state(
 get_job_method=pipeline_client.get_training_pipeline,
 name=shared_state["training_pipeline_name"],
+timeout=timeout,
 )
 
 # Delete the training pipeline

@@ -1,5 +1,7 @@
+import collections
 import re
 import time
+from timeit import default_timer as timer
 
 from typing import Callable
 
@@ -63,3 +65,21 @@ def wait_for_job_state(
 "\nTry increasing the timeout in sample test"
 f"\nLast recorded state: {response.state}"
 )
+
+
+def flaky_test_diagnostic(file_name, test_name, N=20):
+
+import pytest
+
+timing_dict = collections.defaultdict(list)
+for ri in range(N):
+start = timer()
+result = pytest.main(['-s', f'{file_name}::{test_name}'])
+end = timer()
+delta = end-start
+if result == pytest.ExitCode.OK:
+timing_dict['SUCCESS'].append(delta)
+else:
+timing_dict['FAILURE'].append(delta)
+
+return timing_dict

@@ -34,6 +34,9 @@ def teardown(teardown_training_pipeline):
 
 def test_ucaip_generated_create_training_pipeline_sample(capsys, shared_state):
 
+# The return of the cancellation can be flaky; max of 20 runs was 215 sec.
+shared_state["cancel_batch_prediction_job_timeout"] = 300
+
 create_training_pipeline_tabular_regression_sample.create_training_pipeline_tabular_regression_sample(
 project=PROJECT_ID,
 display_name=DISPLAY_NAME,

@@ -0,0 +1,22 @@
+import os
+import pathlib
+import sys
+
+# Run test in snippets directory:
+dir_of_curr_file = os.path.dirname(__file__)
+helper_filepath = pathlib.Path(dir_of_curr_file).parent / 'samples' / 'snippets'
+sys.path.append(helper_filepath.resolve().as_posix())
+os.chdir(helper_filepath.resolve())
+from helpers import flaky_test_diagnostic
+
+# Settings:
+file_name = 'pipeline_service/create_training_pipeline_tabular_regression_sample_test.py'
+test_name = 'test_ucaip_generated_create_training_pipeline_sample'
+timing_dict = flaky_test_diagnostic(file_name, test_name, N=1)
+
+for key, delta_list in timing_dict.items():
+mean_time = sum(delta_list)/len(delta_list)
+max_time = max(delta_list)
+min_time = min(delta_list)
+report_string = f'Result: {key}, mean={mean_time:3.2f}, min={min_time:3.2f}, max={max_time:3.2f}, count={len(delta_list)}'
+print(report_string)