Merge pull request #12905 from sambvfx/artifact-endpoint

[BEAM-8660] Use PortableOptions.artifact_endpoint if provided over value reported from jobserver

Author: Kyle Weaver

sdks/python/apache_beam/runners/portability/local_job_service_test.py

@@ -22,17 +22,22 @@
 import logging
 import unittest
 
+import mock
+
+from apache_beam.options import pipeline_options
 from apache_beam.portability.api import beam_job_api_pb2
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.runners.portability import local_job_service
 from apache_beam.runners.portability.portable_runner import JobServiceHandle
+from apache_beam.runners.portability.portable_runner import PortableRunner
 
 
 class TestJobServicePlan(JobServiceHandle):
   def __init__(self, job_service):
     self.job_service = job_service
     self.options = None
     self.timeout = None
+    self.artifact_endpoint = None
 
   def get_pipeline_options(self):
     return None
@@ -102,6 +107,32 @@ def test_error_messages_after_pipeline_failure(self):
             for m in message_results),
         messages_again)
 
+  def test_artifact_service_override(self):
+    job_service = local_job_service.LocalJobServicer()
+    port = job_service.start_grpc_server()
+
+    test_artifact_endpoint = 'testartifactendpoint:4242'
+
+    options = pipeline_options.PipelineOptions([
+        '--job_endpoint',
+        'localhost:%d' % port,
+        '--artifact_endpoint',
+        test_artifact_endpoint,
+    ])
+    runner = PortableRunner()
+    job_service_handle = runner.create_job_service(options)
+
+    with mock.patch.object(job_service_handle, 'stage') as mocked_stage:
+      job_service_handle.submit(beam_runner_api_pb2.Pipeline())
+      mocked_stage.assert_called_once_with(
+          mock.ANY, test_artifact_endpoint, mock.ANY)
+
+    # Confirm the artifact_endpoint is in the options protobuf
+    options_proto = job_service_handle.get_pipeline_options()
+    self.assertEqual(
+        options_proto['beam:option:artifact_endpoint:v1'],
+        test_artifact_endpoint)
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)

sdks/python/apache_beam/runners/portability/portable_runner.py

@@ -96,6 +96,7 @@ def __init__(self, job_service, options, retain_unknown_options=False):
     self.job_service = job_service
     self.options = options
     self.timeout = options.view_as(PortableOptions).job_server_timeout
+    self.artifact_endpoint = options.view_as(PortableOptions).artifact_endpoint
     self._retain_unknown_options = retain_unknown_options
 
   def submit(self, proto_pipeline):
@@ -105,9 +106,12 @@ def submit(self, proto_pipeline):
     Submit and run the pipeline defined by `proto_pipeline`.
     """
     prepare_response = self.prepare(proto_pipeline)
+    artifact_endpoint = (
+        self.artifact_endpoint or
+        prepare_response.artifact_staging_endpoint.url)
     self.stage(
         proto_pipeline,
-        prepare_response.artifact_staging_endpoint.url,
+        artifact_endpoint,
         prepare_response.staging_session_token)
     return self.run(prepare_response.preparation_id)