[Python] Allow users to pass service name for profiler

CHANGES.md

@@ -65,6 +65,7 @@
 ## New Features / Improvements
 
 * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* Allow passing service name for google-cloud-profiler (Python) ([#26280](https://github.com/apache/beam/issues/26280)).
 
 ## Breaking Changes
 

sdks/python/apache_beam/options/pipeline_options.py

@@ -876,6 +876,16 @@ def validate(self, validator):
 
     return errors
 
+  def lookup_dataflow_service_option(self, key, default=None):
+    if not self.dataflow_service_options:
+      return None
+    elif key in self.dataflow_service_options:
+      return default
+    for service_name in self.dataflow_service_options:
+      if service_name.startswith(key + '='):
+        return service_name.split('=', 1)[1]
+    return None
+
 
 class AzureOptions(PipelineOptions):
   """Azure Blob Storage options."""

sdks/python/apache_beam/runners/worker/data_sampler.py

@@ -29,7 +29,6 @@
 from typing import Iterable
 from typing import List
 from typing import Optional
-from typing import Tuple
 from typing import Union
 
 from apache_beam.coders.coder_impl import CoderImpl

sdks/python/apache_beam/runners/worker/sdk_worker_main.py

@@ -175,20 +175,20 @@ def create_harness(environment, dry_run=False):
 
 def main(unused_argv):
   """Main entry point for SDK Fn Harness."""
-  fn_log_handler, sdk_harness, sdk_pipeline_options = create_harness(os.environ)
-  experiments = sdk_pipeline_options.view_as(DebugOptions).experiments or []
-  dataflow_service_options = (
-      sdk_pipeline_options.view_as(GoogleCloudOptions).dataflow_service_options
-      or [])
-  if (_ENABLE_GOOGLE_CLOUD_PROFILER in experiments) or (
-      _ENABLE_GOOGLE_CLOUD_PROFILER in dataflow_service_options):
+  (fn_log_handler, sdk_harness,
+   sdk_pipeline_options) = create_harness(os.environ)
+  experiments = (sdk_pipeline_options.view_as(DebugOptions).experiments or [])
+  service_name = sdk_pipeline_options.view_as(
+      GoogleCloudOptions).lookup_dataflow_service_option(
+          _ENABLE_GOOGLE_CLOUD_PROFILER, default=os.environ["JOB_NAME"])
+
+  if ((_ENABLE_GOOGLE_CLOUD_PROFILER in experiments) or service_name):
     try:
       import googlecloudprofiler
-      job_id = os.environ["JOB_ID"]
-      job_name = os.environ["JOB_NAME"]
-      if job_id and job_name:
+      service_version = os.environ["JOB_ID"]
+      if service_version and service_name:
         googlecloudprofiler.start(
-            service=job_name, service_version=job_id, verbose=1)
+            service=service_name, service_version=service_version, verbose=1)
         _LOGGER.info('Turning on Google Cloud Profiler.')
       else:
         raise RuntimeError('Unable to find the job id or job name from envvar.')
@@ -199,6 +199,7 @@ def main(unused_argv):
           'https://cloud.google.com/dataflow/docs/guides/profiling-a-pipeline.'
           'For troubleshooting tips with Cloud Profiler see '
           'https://cloud.google.com/profiler/docs/troubleshooting.' % e)
+
   try:
     _LOGGER.info('Python sdk harness starting.')
     sdk_harness.run()