Fix typo in use_single_core_per_container logic. (#22318)

This fixes #22317
apache · Jul 18, 2022 · 70f26ad · 70f26ad
1 parent 1805382
commit 70f26ad
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py
@@ -298,7 +298,7 @@ def __init__(
       container_image = dataflow.SdkHarnessContainerImage()
       container_image.containerImage = container_image_url
       container_image.useSingleCorePerContainer = (
-          common_urns.protocols.MULTI_CORE_BUNDLE_PROCESSING not in
+          common_urns.protocols.MULTI_CORE_BUNDLE_PROCESSING.urn not in
           environment.capabilities)
       container_image.environmentId = id
       for capability in environment.capabilities:

diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py
@@ -187,6 +187,8 @@ def _verify_sdk_harness_container_images_get_set(self, pipeline_options):
         payload=(
             beam_runner_api_pb2.DockerPayload(
                 container_image='dummy_image')).SerializeToString())
+    dummy_env.capabilities.append(
+        common_urns.protocols.MULTI_CORE_BUNDLE_PROCESSING.urn)
     proto_pipeline.components.environments['dummy_env_id'].CopyFrom(dummy_env)
 
     dummy_transform = beam_runner_api_pb2.PTransform(
@@ -203,6 +205,12 @@ def _verify_sdk_harness_container_images_get_set(self, pipeline_options):
     worker_pool = env.proto.workerPools[0]
 
     self.assertEqual(2, len(worker_pool.sdkHarnessContainerImages))
+    # Only one of the environments is missing MULTI_CORE_BUNDLE_PROCESSING.
+    self.assertEqual(
+        1,
+        sum(
+            c.useSingleCorePerContainer
+            for c in worker_pool.sdkHarnessContainerImages))
 
     env_and_image = [(item.environmentId, item.containerImage)
                      for item in worker_pool.sdkHarnessContainerImages]