Support Environment Sensitive Staging

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

@@ -1576,7 +1576,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(
@@ -1585,11 +1585,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',
@@ -1856,7 +1956,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(
@@ -1865,7 +1965,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(), " "))