Pin flaky Python unit tests to FnApiRunner to avoid Prism gRPC deadline timeouts

Author: aIbrahiim

sdks/python/apache_beam/io/gcp/bigquery_test.py

@@ -716,14 +716,16 @@ def test_query_job_exception(self, exception_type, error_message):
   ])
   def test_read_export_exception(self, exception_type, error_message):
 
+    # Mocked BigQuery + Prism can flake in CI; use in-process DirectRunner path.
     with mock.patch.object(beam.io.gcp.bigquery._CustomBigQuerySource,
                            'estimate_size') as mock_estimate,\
       mock.patch.object(bigquery_v2_client.BigqueryV2.TablesService, 'Get'),\
       mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService,
                         'Insert') as mock_query_job, \
       mock.patch('time.sleep'), \
       self.assertRaises(Exception) as exc,\
-      beam.Pipeline() as p:
+      beam.Pipeline(options=PipelineOptions(
+          [], direct_runner_use_prism=False)) as p:
 
       mock_estimate.return_value = None
       mock_query_job.side_effect = exception_type(error_message)

sdks/python/apache_beam/options/pipeline_options.py

@@ -996,6 +996,15 @@ def _add_argparse_args(cls, parser):
         '(in seconds) at which the split requests will be sent, and '
         'fractions is a corresponding list of floating points to use in the '
         'split requests themselves.')
+    parser.add_argument(
+        '--no_direct_runner_use_prism',
+        action='store_false',
+        dest='direct_runner_use_prism',
+        default=True,
+        help='SwitchingDirectRunner will not execute supported batch pipelines '
+        'on PrismRunner (avoids the portable/gRPC subprocess path). Prefer this '
+        'when Prism is slow or flaky, e.g. in constrained or highly parallel '
+        'test environments.')
 
 
 class GoogleCloudOptions(PipelineOptions):

sdks/python/apache_beam/runners/direct/direct_runner.py

@@ -205,7 +205,10 @@ def visit_transform(self, applied_ptransform):
 
     # Check whether all transforms used in the pipeline are supported by the
     # PrismRunner
-    if _PrismRunnerSupportVisitor().accept(pipeline, self._is_interactive):
+    direct_opts = options.view_as(DirectOptions)
+    use_prism = getattr(direct_opts, 'direct_runner_use_prism', True)
+    if (use_prism and
+        _PrismRunnerSupportVisitor().accept(pipeline, self._is_interactive)):
       _LOGGER.info('Running pipeline with PrismRunner.')
       from apache_beam.runners.portability import prism_runner
       runner = prism_runner.PrismRunner()

sdks/python/apache_beam/yaml/yaml_transform_unit_test.py

@@ -50,9 +50,13 @@
 
 
 def new_pipeline():
+  # Default DirectRunner may run batch workloads on Prism (gRPC subprocess);
+  # disable for stable unit tests (DirectOptions.direct_runner_use_prism).
   return beam.Pipeline(
       options=beam.options.pipeline_options.PipelineOptions(
-          pickle_library='cloudpickle'))
+          [],
+          pickle_library='cloudpickle',
+          direct_runner_use_prism=False))
 
 
 @unittest.skipIf(jsonschema is None, "Yaml dependencies not installed")
@@ -652,7 +656,7 @@ def test_preprocess_windowing_custom_type(self):
             size: 4
           input: {{input: input}}
       output: {result['transforms'][0]['__uuid__']}
-      config: 
+      config:
         error_handling: {{}}
     '''
     self.assertYaml(expected, result)
@@ -775,7 +779,7 @@ def test_preprocess_windowing_other_type_with_no_inputs(self):
             type: fixed
             size: 4
       output: {result['transforms'][1]["__uuid__"]}
-      config: 
+      config:
         error_handling: {{}}
     '''
     self.maxDiff = 1e9