Add an experiment to make prism server a singleton (apache#34623)

shunping · web-flow · commit bebfcbc84bd9 · 2025-04-15T15:25:24.000-04:00
* Add an experiment flag to make prism job server a singleton.

* Attempt to fix the failed tests.

* Replace singleton code with util.shared.Shared.

* Fix lints
diff --git a/sdks/python/apache_beam/runners/portability/prism_runner.py b/sdks/python/apache_beam/runners/portability/prism_runner.py
@@ -39,6 +39,7 @@
 from apache_beam.runners.portability import job_server
 from apache_beam.runners.portability import portable_runner
 from apache_beam.transforms import environments
+from apache_beam.utils import shared
 from apache_beam.utils import subprocess_server
 from apache_beam.version import __version__ as beam_version
 
@@ -56,6 +57,8 @@ class PrismRunner(portable_runner.PortableRunner):
   """A runner for launching jobs on Prism, automatically downloading and
   starting a Prism instance if needed.
   """
+  shared_handle = shared.Shared()
+
   def default_environment(
       self,
       options: pipeline_options.PipelineOptions) -> environments.Environment:
@@ -66,7 +69,12 @@ def default_environment(
     return super().default_environment(options)
 
   def default_job_server(self, options):
-    return job_server.StopOnExitJobServer(PrismJobServer(options))
+    debug_options = options.view_as(pipeline_options.DebugOptions)
+    get_job_server = lambda: job_server.StopOnExitJobServer(
+        PrismJobServer(options))
+    if debug_options.lookup_experiment("enable_prism_server_singleton"):
+      return PrismRunner.shared_handle.acquire(get_job_server)
+    return get_job_server()
 
   def create_job_service_handle(self, job_service, options):
     return portable_runner.JobServiceHandle(
diff --git a/sdks/python/apache_beam/runners/portability/prism_runner_test.py b/sdks/python/apache_beam/runners/portability/prism_runner_test.py
@@ -39,6 +39,7 @@
 from apache_beam.runners.portability import prism_runner
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
+from apache_beam.utils import shared
 
 # Run as
 #
@@ -381,6 +382,38 @@ def test_with_remote_path(self, has_cache_bin, has_cache_zip, ignore_cache):
             mock_zipfile_init.assert_called_once()
 
 
+class PrismRunnerSingletonTest(unittest.TestCase):
+  @parameterized.expand([True, False])
+  def test_singleton(self, enable_singleton):
+    if enable_singleton:
+      options = DebugOptions(["--experiment=enable_prism_server_singleton"])
+    else:
+      options = DebugOptions()
+
+    runner = prism_runner.PrismRunner()
+    with mock.patch(
+        'apache_beam.runners.portability.prism_runner.PrismJobServer'
+    ) as mock_prism_server:
+
+      # Reset the class-level singleton for every fresh run
+      prism_runner.PrismRunner.shared_handle = shared.Shared()
+
+      runner = prism_runner.PrismRunner()
+      runner.default_job_server(options)
+
+      mock_prism_server.assert_called_once()
+      mock_prism_server.reset_mock()
+
+      runner = prism_runner.PrismRunner()
+      runner.default_job_server(options)
+      if enable_singleton:
+        # If singleton is enabled, we won't try to create a new server for the
+        # second run.
+        mock_prism_server.assert_not_called()
+      else:
+        mock_prism_server.assert_called_once()
+
+
 if __name__ == '__main__':
   # Run the tests.
   logging.getLogger().setLevel(logging.INFO)
