Add an experiment flag to make prism job server a singleton.

Author: shunping

Diff for: sdks/python/apache_beam/runners/portability/prism_runner.py

@@ -28,6 +28,7 @@
 import shutil
 import stat
 import subprocess
+import threading
 import typing
 import urllib
 import zipfile
@@ -56,6 +57,20 @@ class PrismRunner(portable_runner.PortableRunner):
   """A runner for launching jobs on Prism, automatically downloading and
   starting a Prism instance if needed.
   """
+  __singleton = None
+  __singleton_lock = threading.Lock()
+
+  @staticmethod
+  def get_job_server(options):
+    debug_options = options.view_as(pipeline_options.DebugOptions)
+    if debug_options.lookup_experiment("enable_prism_server_singleton"):
+      with PrismRunner.__singleton_lock:
+        if PrismRunner.__singleton is None:
+          PrismRunner.__singleton = PrismJobServer(options)
+      return PrismRunner.__singleton
+
+    return PrismJobServer(options)
+
   def default_environment(
       self,
       options: pipeline_options.PipelineOptions) -> environments.Environment:
@@ -66,7 +81,7 @@ def default_environment(
     return super().default_environment(options)
 
   def default_job_server(self, options):
-    return job_server.StopOnExitJobServer(PrismJobServer(options))
+    return job_server.StopOnExitJobServer(PrismRunner.get_job_server(options))
 
   def create_job_service_handle(self, job_service, options):
     return portable_runner.JobServiceHandle(
@@ -92,6 +107,22 @@ def __init__(self, options):
 
     job_options = options.view_as(pipeline_options.JobServerOptions)
     self._job_port = job_options.job_port
+    self._lock = threading.Lock()
+    self._started = False
+    self._endpoint = None
+
+  def start(self):
+    with self._lock:
+      if not self._started:
+        self._endpoint = super().start()
+        self._started = True
+    return self._endpoint
+
+  def stop(self):
+    with self._lock:
+      if self._started:
+        super().stop()
+        self._started = False
 
   @classmethod
   def maybe_unzip_and_make_executable(

Diff for: sdks/python/apache_beam/runners/portability/prism_runner_test.py

@@ -381,6 +381,47 @@ 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()
+
+    with mock.patch(
+        'apache_beam.runners.portability.job_server.subprocess_server.SubprocessServer.start'  # pylint: disable=line-too-long
+    ) as mock_start:
+      # Reset the class-level singleton for every fresh run
+      prism_runner.PrismRunner._PrismRunner__singleton = None
+
+      try:
+        with beam.Pipeline(options=options,
+                           runner=prism_runner.PrismRunner()) as p:
+          _ = p | "Create Elements" >> beam.Create(
+              range(5)) | "Squares" >> beam.Map(lambda x: x**2)
+      except:  # pylint: disable=bare-except
+        pass
+
+      mock_start.assert_called_once()
+      mock_start.reset_mock()
+
+      try:
+        with beam.Pipeline(options=options,
+                           runner=prism_runner.PrismRunner()) as p:
+          _ = p | "Create Elements" >> beam.Create(
+              range(5)) | "Squares" >> beam.Map(lambda x: x**2)
+      except:  # pylint: disable=bare-except
+        pass
+
+      if enable_singleton:
+        # If singleton is enabled, we won't try to start a new server for the
+        # second run.
+        mock_start.assert_not_called()
+      else:
+        mock_start.assert_called_once()
+
+
 if __name__ == '__main__':
   # Run the tests.
   logging.getLogger().setLevel(logging.INFO)