More eagerly merge compatible environments.

Specifically, merge environments that are identical except for their
dependencies, but where one environment's dependencies are a superset
of the others.

This will reduce the number of distinct environments of the same
type, and the associated materialization costs of moving between
them.

Ideally this would be done in a runner for environments that are
already co-located on the same worker, but that is a further
optimization.

Author: robertwb

sdks/python/apache_beam/runners/common.py

@@ -2008,43 +2008,65 @@ def validate_transform(transform_id):
     validate_transform(t)
 
 
-def merge_common_environments(pipeline_proto, inplace=False):
-  def dep_key(dep):
-    if dep.type_urn == common_urns.artifact_types.FILE.urn:
-      payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
-          dep.type_payload)
-      if payload.sha256:
-        type_info = 'sha256', payload.sha256
-      else:
-        type_info = 'path', payload.path
-    elif dep.type_urn == common_urns.artifact_types.URL.urn:
-      payload = beam_runner_api_pb2.ArtifactUrlPayload.FromString(
-          dep.type_payload)
-      if payload.sha256:
-        type_info = 'sha256', payload.sha256
-      else:
-        type_info = 'url', payload.url
+def _dep_key(dep):
+  if dep.type_urn == common_urns.artifact_types.FILE.urn:
+    payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
+        dep.type_payload)
+    if payload.sha256:
+      type_info = 'sha256', payload.sha256
     else:
-      type_info = dep.type_urn, dep.type_payload
-    return type_info, dep.role_urn, dep.role_payload
-
-  def base_env_key(env):
-    return (
-        env.urn,
-        env.payload,
-        tuple(sorted(env.capabilities)),
-        tuple(sorted(env.resource_hints.items())),
-        tuple(sorted(dep_key(dep) for dep in env.dependencies)))
-
-  def env_key(env):
-    return tuple(
-        sorted(
-            base_env_key(e)
-            for e in environments.expand_anyof_environments(env)))
+      type_info = 'path', payload.path
+  elif dep.type_urn == common_urns.artifact_types.URL.urn:
+    payload = beam_runner_api_pb2.ArtifactUrlPayload.FromString(
+        dep.type_payload)
+    if payload.sha256:
+      type_info = 'sha256', payload.sha256
+    else:
+      type_info = 'url', payload.url
+  else:
+    type_info = dep.type_urn, dep.type_payload
+  return type_info, dep.role_urn, dep.role_payload
+
+
+def _expanded_dep_keys(dep):
+  if (dep.type_urn == common_urns.artifact_types.FILE.urn and
+      dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn):
+    payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
+        dep.type_payload)
+    role = beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString(
+        dep.role_payload)
+    if role.staged_name == 'submission_environment_dependencies.txt':
+      return
+    elif role.staged_name == 'requirements.txt':
+      with open(payload.path) as fin:
+        for line in fin:
+          yield 'requirements.txt', line.strip()
+      return
+
+  yield _dep_key(dep)
+
+
+def _base_env_key(env, include_deps=True):
+  return (
+      env.urn,
+      env.payload,
+      tuple(sorted(env.capabilities)),
+      tuple(sorted(env.resource_hints.items())),
+      tuple(sorted(_dep_key(dep)
+                   for dep in env.dependencies)) if include_deps else None)
+
 
+def _env_key(env):
+  return tuple(
+      sorted(
+          _base_env_key(e)
+          for e in environments.expand_anyof_environments(env)))
+
+
+def merge_common_environments(pipeline_proto, inplace=False):
   canonical_environments = collections.defaultdict(list)
   for env_id, env in pipeline_proto.components.environments.items():
-    canonical_environments[env_key(env)].append(env_id)
+    canonical_environments[_env_key(env)].append(env_id)
 
   if len(canonical_environments) == len(pipeline_proto.components.environments):
     # All environments are already sufficiently distinct.
@@ -2055,6 +2077,55 @@ def env_key(env):
       for es in canonical_environments.values() for e in es
   }
 
+  return update_environments(pipeline_proto, environment_remappings, inplace)
+
+
+def merge_superset_dep_environments(pipeline_proto):
+  """Merges all environemnts A and B where A and B are equivalent except that
+  A has a superset of the dependencies of B.
+  """
+  docker_envs = {}
+  for env_id, env in pipeline_proto.components.environments.items():
+    docker_env = environments.resolve_anyof_environment(
+        env, common_urns.environments.DOCKER.urn)
+    if docker_env.urn == common_urns.environments.DOCKER.urn:
+      docker_envs[env_id] = docker_env
+
+  has_base_and_dep = collections.defaultdict(set)
+  env_scores = {
+      env_id: (len(env.dependencies), env_id)
+      for (env_id, env) in docker_envs.items()
+  }
+
+  for env_id, env in docker_envs.items():
+    base_key = _base_env_key(env, include_deps=False)
+    has_base_and_dep[base_key, None].add(env_id)
+    for dep in env.dependencies:
+      for dep_key in _expanded_dep_keys(dep):
+        has_base_and_dep[base_key, dep_key].add(env_id)
+
+  environment_remappings = {}
+  for env_id, env in docker_envs.items():
+    base_key = _base_env_key(env, include_deps=False)
+    # This is the set of all environments that have at least all of env's deps.
+    candidates = set.intersection(
+        has_base_and_dep[base_key, None],
+        *[
+            has_base_and_dep[base_key, dep_key] for dep in env.dependencies
+            for dep_key in _expanded_dep_keys(dep)
+        ])
+    # Choose the maximal one.
+    best = max(candidates, key=env_scores.get)
+    if best != env_id:
+      environment_remappings[env_id] = best
+
+  return update_environments(pipeline_proto, environment_remappings)
+
+
+def update_environments(pipeline_proto, environment_remappings, inplace=False):
+  if not environment_remappings:
+    return pipeline_proto
+
   if not inplace:
     pipeline_proto = copy.copy(pipeline_proto)
 
@@ -2063,16 +2134,16 @@ def env_key(env):
       # TODO(https://github.com/apache/beam/issues/30876): Remove this
       #  workaround.
       continue
-    if t.environment_id:
+    if t.environment_id and t.environment_id in environment_remappings:
       t.environment_id = environment_remappings[t.environment_id]
   for w in pipeline_proto.components.windowing_strategies.values():
     if w.environment_id not in pipeline_proto.components.environments:
       # TODO(https://github.com/apache/beam/issues/30876): Remove this
       #  workaround.
       continue
-    if w.environment_id:
+    if w.environment_id and w.environment_id in environment_remappings:
       w.environment_id = environment_remappings[w.environment_id]
-  for e in set(pipeline_proto.components.environments.keys()) - set(
+  for e in set(environment_remappings.keys()) - set(
       environment_remappings.values()):
     del pipeline_proto.components.environments[e]
   return pipeline_proto

sdks/python/apache_beam/runners/common_test.py

@@ -17,6 +17,10 @@
 
 # pytype: skip-file
 
+import hashlib
+import os
+import random
+import tempfile
 import unittest
 
 import hamcrest as hc
@@ -26,10 +30,12 @@
 from apache_beam.io.restriction_trackers import OffsetRestrictionTracker
 from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
 from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.portability import common_urns
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.runners.common import DoFnSignature
 from apache_beam.runners.common import PerWindowInvoker
 from apache_beam.runners.common import merge_common_environments
+from apache_beam.runners.common import merge_superset_dep_environments
 from apache_beam.runners.portability.expansion_service_test import FibTransform
 from apache_beam.runners.sdf_utils import SplitResultPrimary
 from apache_beam.runners.sdf_utils import SplitResultResidual
@@ -632,6 +638,111 @@ def test_equal_environments_merged(self):
                 pipeline_proto.components.windowing_strategies.values())),
         1)
 
+  def _make_dep(self, path):
+    hasher = hashlib.sha256()
+    if os.path.exists(path):
+      with open(path, 'rb') as fin:
+        hasher.update(fin.read())
+    else:
+      # A fake file, identified only by its path.
+      hasher.update(path.encode('utf-8'))
+    return beam_runner_api_pb2.ArtifactInformation(
+        type_urn=common_urns.artifact_types.FILE.urn,
+        type_payload=beam_runner_api_pb2.ArtifactFilePayload(
+            path=path, sha256=hasher.hexdigest()).SerializeToString(),
+        role_urn=common_urns.artifact_roles.STAGING_TO.urn,
+        role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload(
+            staged_name=os.path.basename(path)).SerializeToString())
+
+  def _docker_env(self, id, deps=()):
+    return beam_runner_api_pb2.Environment(
+        urn=common_urns.environments.DOCKER.urn,
+        payload=id.encode('utf8'),
+        dependencies=[self._make_dep(path) for path in deps],
+    )
+
+  def test_subset_deps_environments_merged(self):
+    environments = {
+        'A': self._docker_env('A'),
+        'Ax': self._docker_env('A', ['x']),
+        'Ay': self._docker_env('A', ['y']),
+        'Axy': self._docker_env('A', ['x', 'y']),
+        'Bx': self._docker_env('B', ['x']),
+        'Bxy': self._docker_env('B', ['x', 'y']),
+        'Byz': self._docker_env('B', ['y', 'z']),
+    }
+    transforms = {
+        env_id: beam_runner_api_pb2.PTransform(
+            unique_name=env_id, environment_id=env_id)
+        for env_id in environments.keys()
+    }
+    pipeline_proto = merge_superset_dep_environments(
+        beam_runner_api_pb2.Pipeline(
+            components=beam_runner_api_pb2.Components(
+                environments=environments, transforms=transforms)))
+
+    # These can all be merged into the same environment.
+    self.assertEqual(
+        pipeline_proto.components.transforms['A'].environment_id, 'Axy')
+    self.assertEqual(
+        pipeline_proto.components.transforms['Ax'].environment_id, 'Axy')
+    self.assertEqual(
+        pipeline_proto.components.transforms['Ay'].environment_id, 'Axy')
+    self.assertEqual(
+        pipeline_proto.components.transforms['Axy'].environment_id, 'Axy')
+    # Despite having the same dependencies, these must be merged into their own.
+    self.assertEqual(
+        pipeline_proto.components.transforms['Bx'].environment_id, 'Bxy')
+    self.assertEqual(
+        pipeline_proto.components.transforms['Bxy'].environment_id, 'Bxy')
+    # This is not a subset of any, must be left alone.
+    self.assertEqual(
+        pipeline_proto.components.transforms['Byz'].environment_id, 'Byz')
+
+  def test_subset_deps_environments_merged_with_requirements_txt(self):
+    with tempfile.TemporaryDirectory() as tmpdir:
+
+      def make_file(basename, content):
+        subdir = tempfile.TemporaryDirectory(dir=tmpdir, delete=False).name
+        path = os.path.join(subdir, basename)
+        with open(path, 'w') as fout:
+          fout.write(content)
+        return path
+
+      def make_py_deps(*pkgs):
+        return [
+            make_file('requirements.txt', '\n'.join(pkgs)),
+            make_file(
+                'submission_environment_dependencies.txt', str(
+                    random.random())),
+        ] + [make_file(pkg, pkg) for pkg in pkgs]
+
+      environments = {
+          'A': self._docker_env('A'),
+          'Ax': self._docker_env('A', make_py_deps('x')),
+          'Ay': self._docker_env('A', make_py_deps('y')),
+          'Axy': self._docker_env('A', make_py_deps('x', 'y')),
+      }
+      transforms = {
+          env_id: beam_runner_api_pb2.PTransform(
+              unique_name=env_id, environment_id=env_id)
+          for env_id in environments.keys()
+      }
+      pipeline_proto = merge_superset_dep_environments(
+          beam_runner_api_pb2.Pipeline(
+              components=beam_runner_api_pb2.Components(
+                  environments=environments, transforms=transforms)))
+
+      # These can all be merged into the same environment.
+      self.assertEqual(
+          pipeline_proto.components.transforms['A'].environment_id, 'Axy')
+      self.assertEqual(
+          pipeline_proto.components.transforms['Ax'].environment_id, 'Axy')
+      self.assertEqual(
+          pipeline_proto.components.transforms['Ay'].environment_id, 'Axy')
+      self.assertEqual(
+          pipeline_proto.components.transforms['Axy'].environment_id, 'Axy')
+
   def test_external_merged(self):
     p = beam.Pipeline()
     # This transform recursively creates several external environments.

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -45,6 +45,7 @@
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.runners.common import group_by_key_input_visitor
 from apache_beam.runners.common import merge_common_environments
+from apache_beam.runners.common import merge_superset_dep_environments
 from apache_beam.runners.dataflow.internal.clients import dataflow as dataflow_api
 from apache_beam.runners.runner import PipelineResult
 from apache_beam.runners.runner import PipelineRunner
@@ -434,7 +435,8 @@ def run_pipeline(self, pipeline, options, pipeline_proto=None):
       self.proto_pipeline.components.environments[env_id].CopyFrom(
           environments.resolve_anyof_environment(
               env, common_urns.environments.DOCKER.urn))
-    self.proto_pipeline = merge_common_environments(self.proto_pipeline)
+    self.proto_pipeline = merge_common_environments(
+        merge_superset_dep_environments(self.proto_pipeline))
 
     # Optimize the pipeline if it not streaming and the pre_optimize
     # experiment is set.

sdks/python/apache_beam/runners/portability/fn_api_runner/fn_runner.py

@@ -64,6 +64,7 @@
 from apache_beam.runners import runner
 from apache_beam.runners.common import group_by_key_input_visitor
 from apache_beam.runners.common import merge_common_environments
+from apache_beam.runners.common import merge_superset_dep_environments
 from apache_beam.runners.common import validate_pipeline_graph
 from apache_beam.runners.portability import portable_metrics
 from apache_beam.runners.portability.fn_api_runner import execution
@@ -216,7 +217,8 @@ def run_via_runner_api(
     if direct_options.direct_embed_docker_python:
       pipeline_proto = self.embed_default_docker_image(pipeline_proto)
     pipeline_proto = merge_common_environments(
-        self.resolve_any_environments(pipeline_proto))
+        self.resolve_any_environments(
+          merge_superset_dep_environments(pipeline_proto)))
     stage_context, stages = self.create_stages(pipeline_proto)
     return self.run_stages(stage_context, stages)