Support Environment Sensitive Staging

CHANGES.md

@@ -129,6 +129,7 @@
 * (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)).
 * (Python) `ConcatSource` no longer raises `ZeroDivisionError` when every source bundle reports a weight of 0, e.g. bundles smaller than 1 MB ([#21316](https://github.com/apache/beam/issues/21316)).
 * (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)).
+* (Python) Dataflow: artifacts staged by different environments under the same name (e.g. requirements.txt from the main pipeline and a Beam YAML pythonPackage provider) no longer overwrite each other in the staging location, and the Python SDK container stages each environment under its own directory. Previously the losing environment failed the staged-file integrity check and crash-looped, or silently installed another environment's dependencies ([#40699](https://github.com/apache/beam/issues/40699)).
 * (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#40644](https://github.com/apache/beam/issues/40644)).
 * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)).

sdks/python/apache_beam/runners/dataflow/internal/apiclient.py

@@ -599,6 +599,8 @@ def _stage_resources(self, pipeline, options):
     resources = []
     staged_paths = {}
     staged_hashes = {}
+    # remote_name -> sha256, to detect different content staged under one name.
+    staged_names = {}
     for _, env in sorted(pipeline.components.environments.items(),
                          key=lambda kv: kv[0]):
       for dep in env.dependencies:
@@ -611,10 +613,8 @@ def _stage_resources(self, pipeline, options):
           remote_name = (
               beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString(
                   dep.role_payload)).staged_name
-          is_staged_role = True
         else:
           remote_name = os.path.basename(type_payload.path)
-          is_staged_role = False
         # compute sha256 even if caching is disabled.
         # This is used to check the payload integrity along with caching.
         if not type_payload.sha256:
@@ -626,31 +626,31 @@ def _stage_resources(self, pipeline, options):
               type_payload.path,
               type_payload.sha256)
           remote_name = staged_hashes[type_payload.sha256]
-          if is_staged_role:
-            # We should not be overriding this, as dep.role_payload.staged_name
-            # refers to the desired name on the worker, whereas staged_name
-            # refers to its placement in a distributed filesystem.
-            # TODO(heejong): Clean this up.
-            dep.role_payload = beam_runner_api_pb2.ArtifactStagingToRolePayload(
-                staged_name=remote_name).SerializeToString()
         elif type_payload.path and type_payload.path in staged_paths:
           _LOGGER.info(
               'Found duplicated artifact path: %s (%s)',
               type_payload.path,
               type_payload.sha256)
           remote_name = staged_paths[type_payload.path]
-          if is_staged_role:
-            # We should not be overriding this, as dep.role_payload.staged_name
-            # refers to the desired name on the worker, whereas staged_name
-            # refers to its placement in a distributed filesystem.
-            # TODO(heejong): Clean this up.
-            dep.role_payload = beam_runner_api_pb2.ArtifactStagingToRolePayload(
-                staged_name=remote_name).SerializeToString()
         else:
+          if staged_names.get(remote_name,
+                              type_payload.sha256) != (type_payload.sha256):
+            # Different content already staged under this name (e.g. each
+            # Python environment stages its own requirements.txt).  Stage this
+            # one as <name>-<sha256><ext> instead of overwriting it.  The
+            # worker-local name (role_payload.staged_name) is unaffected.
+            stem, ext = os.path.splitext(remote_name)
+            remote_name = '%s-%s%s' % (stem, type_payload.sha256, ext)
+            _LOGGER.info(
+                'Artifact %s collides with a staged artifact of the same name '
+                'and different content; staging it as %s.',
+                type_payload.path,
+                remote_name)
           resources.append(
               (type_payload.path, remote_name, type_payload.sha256))
           staged_paths[type_payload.path] = remote_name
           staged_hashes[type_payload.sha256] = remote_name
+          staged_names[remote_name] = type_payload.sha256
 
         if FileSystems.get_scheme(
             google_cloud_options.staging_location) == GCSFileSystem.scheme():

sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py

@@ -1651,7 +1651,7 @@ def test_stage_resources(self):
                             role_urn=common_urns.artifact_roles.STAGING_TO.urn,
                             role_payload=beam_runner_api_pb2.
                             ArtifactStagingToRolePayload(
-                                staged_name='baz1').SerializeToString()),
+                                staged_name='baz2').SerializeToString()),
                         beam_runner_api_pb2.ArtifactInformation(
                             type_urn=common_urns.artifact_types.URL.urn,
                             type_payload=beam_runner_api_pb2.ArtifactUrlPayload(
@@ -1660,11 +1660,111 @@ def test_stage_resources(self):
                             role_urn=common_urns.artifact_roles.STAGING_TO.urn,
                             role_payload=beam_runner_api_pb2.
                             ArtifactStagingToRolePayload(
-                                staged_name='renamed1').SerializeToString())
+                                staged_name='renamed2').SerializeToString())
                     ])
             }))
     self.assertEqual(pipeline, pipeline_expected)
 
+  def _stage_resources_for(self, environments):
+    pipeline_options = PipelineOptions([
+        '--temp_location',
+        'gs://test-location/temp',
+        '--staging_location',
+        'gs://test-location/staging',
+        '--no_auth'
+    ])
+
+    def artifact(path, staged_name, sha256):
+      return beam_runner_api_pb2.ArtifactInformation(
+          type_urn=common_urns.artifact_types.FILE.urn,
+          type_payload=beam_runner_api_pb2.ArtifactFilePayload(
+              path=path, sha256=sha256).SerializeToString(),
+          role_urn=common_urns.artifact_roles.STAGING_TO.urn,
+          role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload(
+              staged_name=staged_name).SerializeToString())
+
+    pipeline = beam_runner_api_pb2.Pipeline(
+        components=beam_runner_api_pb2.Components(
+            environments={
+                env_id: beam_runner_api_pb2.Environment(
+                    dependencies=[artifact(*dep) for dep in deps])
+                for env_id, deps in environments.items()
+            }))
+    client = apiclient.DataflowApplicationClient(pipeline_options)
+    with mock.patch.object(apiclient._LegacyDataflowStager,
+                           'stage_job_resources') as mock_stager:
+      client._stage_resources(pipeline, pipeline_options)
+    staged = mock_stager.call_args[0][0]
+    deps = {
+        env_id: [(
+            beam_runner_api_pb2.ArtifactUrlPayload.FromString(
+                d.type_payload).url,
+            beam_runner_api_pb2.ArtifactUrlPayload.FromString(
+                d.type_payload).sha256,
+            beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString(
+                d.role_payload).staged_name) for d in env.dependencies]
+        for env_id, env in pipeline.components.environments.items()
+    }
+    return staged, deps
+
+  def test_stage_resources_same_name_different_content(self):
+    # Two environments each stage their own requirements.txt (e.g. the main
+    # pipeline and a Beam YAML pythonPackage provider).  The second must not
+    # overwrite the first, and both keep staged_name=requirements.txt so the
+    # SDK container still finds the file.
+    staged, deps = self._stage_resources_for({
+        'env1': [('/tmp/a/requirements.txt', 'requirements.txt', 'aaaa')],
+        'env2': [('/tmp/b/requirements.txt', 'requirements.txt', 'bbbb'),
+                 ('/tmp/b/pkg.tar.gz', 'pkg.tar.gz', 'cccc')],
+    })
+    self.assertEqual(
+        staged,
+        [('/tmp/a/requirements.txt', 'requirements.txt', 'aaaa'),
+         ('/tmp/b/requirements.txt', 'requirements-bbbb.txt', 'bbbb'),
+         ('/tmp/b/pkg.tar.gz', 'pkg.tar.gz', 'cccc')])
+    self.assertEqual(
+        deps['env1'],
+        [(
+            'gs://test-location/staging/requirements.txt',
+            'aaaa',
+            'requirements.txt')])
+    self.assertEqual(
+        deps['env2'],
+        [(
+            'gs://test-location/staging/requirements-bbbb.txt',
+            'bbbb',
+            'requirements.txt'),
+         ('gs://test-location/staging/pkg.tar.gz', 'cccc', 'pkg.tar.gz')])
+
+  def test_stage_resources_same_content_keeps_staged_name(self):
+    # Identical bytes are uploaded once, but each environment keeps the
+    # worker-local name it asked for: role_payload is never rewritten.
+    staged, deps = self._stage_resources_for({
+        'env1': [('/tmp/a/requirements.txt', 'requirements.txt', 'same')],
+        'env2': [('/tmp/b/requirements.txt', 'requirements.txt', 'same'),
+                 ('/tmp/b/other.txt', 'other.txt', 'same')],
+        'env3': [('/tmp/c/requirements.txt', 'requirements.txt', 'diff')],
+        'env4': [('/tmp/d/requirements.txt', 'requirements.txt', 'diff')],
+    })
+    self.assertEqual(
+        staged,
+        [('/tmp/a/requirements.txt', 'requirements.txt', 'same'),
+         ('/tmp/c/requirements.txt', 'requirements-diff.txt', 'diff')])
+    self.assertEqual(
+        deps['env2'],
+        [(
+            'gs://test-location/staging/requirements.txt',
+            'same',
+            'requirements.txt'),
+         ('gs://test-location/staging/requirements.txt', 'same', 'other.txt')])
+    # env4 reuses env3's renamed object but still asks for requirements.txt.
+    self.assertEqual(
+        deps['env4'],
+        [(
+            'gs://test-location/staging/requirements-diff.txt',
+            'diff',
+            'requirements.txt')])
+
   def test_set_dataflow_service_option(self):
     pipeline_options = PipelineOptions([
         '--dataflow_service_option',
@@ -1931,7 +2031,7 @@ def exists_return_value(*args):
                             role_urn=common_urns.artifact_roles.STAGING_TO.urn,
                             role_payload=beam_runner_api_pb2.
                             ArtifactStagingToRolePayload(
-                                staged_name='baz1').SerializeToString()),
+                                staged_name='baz2').SerializeToString()),
                         beam_runner_api_pb2.ArtifactInformation(
                             type_urn=common_urns.artifact_types.URL.urn,
                             type_payload=beam_runner_api_pb2.ArtifactUrlPayload(
@@ -1940,7 +2040,7 @@ def exists_return_value(*args):
                             role_urn=common_urns.artifact_roles.STAGING_TO.urn,
                             role_payload=beam_runner_api_pb2.
                             ArtifactStagingToRolePayload(
-                                staged_name='renamed1').SerializeToString())
+                                staged_name='renamed2').SerializeToString())
                     ])
             }))
     self.assertEqual(pipeline, pipeline_expected)

sdks/python/container/boot.go

@@ -211,7 +211,12 @@ func launchSDKProcess() error {
 		defer cleanupFunc()
 	}
 
-	dir := filepath.Join(*semiPersistDir, "staged")
+	// Stage under a per-container directory
+	// So that, SDK containers for different environments sharing the same
+	// semi-persistent volume cannot overwrite each other's fixed-name
+	// artifacts (requirements.txt, pickled_main_session, etc.,).
+	semiPersist := filepath.Join(*semiPersistDir, *id)
+	dir := filepath.Join(semiPersist, "staged")
 	files, err := artifact.Materialize(ctx, *artifactEndpoint, info.GetDependencies(), info.GetRetrievalToken(), dir)
 	if err != nil {
 		fmtErr := fmt.Errorf("failed to retrieve staged files: %v", err)
@@ -257,7 +262,9 @@ func launchSDKProcess() error {
 		logger.Fatalf(ctx, "Failed to load pipeline options to worker: %v", err)
 	}
 
-	os.Setenv("SEMI_PERSISTENT_DIRECTORY", *semiPersistDir)
+	// Point the SDK worker at the per-environment root so that it keeps
+	// resolving <SEMI_PERSISTENT_DIRECTORY>/staged/... (e.g. pickled_main_session).
+	os.Setenv("SEMI_PERSISTENT_DIRECTORY", semiPersist)
 	os.Setenv("LOGGING_API_SERVICE_DESCRIPTOR", (&pipepb.ApiServiceDescriptor{Url: *loggingEndpoint}).String())
 	os.Setenv("CONTROL_API_SERVICE_DESCRIPTOR", (&pipepb.ApiServiceDescriptor{Url: *controlEndpoint}).String())
 	os.Setenv("RUNNER_CAPABILITIES", strings.Join(info.GetRunnerCapabilities(), " "))