Support Environment Sensitive Staging

.github/trigger_files/beam_PostCommit_Python.json

@@ -1,5 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run.",
-  "pr": "40241",
-  "modification": 66
+  "pr": "40699",
+  "modification": 68
 }

CHANGES.md

@@ -99,6 +99,8 @@
 * (Python) `Regex.all_matches` now returns every group of the regex, instead of failing with a `TypeError` when the regex has no groups and dropping nested or unmatched trailing groups ([#40562](https://github.com/apache/beam/issues/40562)).
 * (Python) `--dataflow_endpoint` now defaults to `''` instead of `https://dataflow.googleapis.com`. When unset, the Google Cloud client library resolves the Dataflow endpoint and honors `GOOGLE_CLOUD_UNIVERSE_DOMAIN`, so Dataflow jobs can be submitted in non-default universes. Code that reads `GoogleCloudOptions.dataflow_endpoint` should handle an empty value ([#40569](https://github.com/apache/beam/issues/40569)).
 * (Python) The `xgboost` extras now install `xgboost>=2.0.0,<4` and `datatable>=1.1.0,<2`, moving both the lower and upper bounds for both. The previous minimum `xgboost` and `datatable` versions (`1.6.0` and `1.0.0` respectively) are no longer allowed ([#40694](https://github.com/apache/beam/issues/40694)).
+* (Python) The Python SDK container now downloads each environment's staged files to `<semi_persist_dir>/<worker id>/staged` instead of `<semi_persist_dir>/staged`, so environments sharing a worker no longer overwrite each other's files. Code that hardcodes a staged file location, such as `/tmp/staged/...` or `/var/opt/google/staged/...`, must instead resolve it on the worker at runtime, for example `os.path.join(os.environ['SEMI_PERSISTENT_DIRECTORY'], 'staged', 'my_file')`. See [Accessing staged files on the workers](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/#accessing-staged-files-on-the-workers) ([#40699](https://github.com/apache/beam/issues/40699)).
+* (Python) Each Python environment of a pipeline (for example a Beam YAML `pythonPackage` provider) now installs its own `requirements.txt`. Previously it could end up installing the main pipeline's requirements instead, so a provider that relied on dependencies listed only in the main pipeline's requirements must now declare them itself ([#40699](https://github.com/apache/beam/issues/40699)).
 
 ## Deprecations
 
@@ -133,6 +135,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) `ReadBigQueryChangeHistory` no longer fails at pipeline construction with `TypeCheckError: The transform 'ExecuteQueries' requires PCollections of type ... _QueryRange` when pipeline type checking is enabled (the default); a regression in 2.76.0 ([#40618](https://github.com/apache/beam/issues/40618)).
 * (Python) `ReadBigQueryChangeHistory` with `change_function='CHANGES'` no longer fails reading its temp table with a Storage Read API `500 InternalServerError`: the `_CHANGE_IS_FOR_UPDATE` pseudo-column is now excluded from the query instead of being written to the temp table ([#40618](https://github.com/apache/beam/issues/40618)).
 * (Python) `ReadBigQueryChangeHistory` now backtick-quotes the column names it puts in the generated CHANGES/APPENDS query (`columns`, `change_type_column`, `change_timestamp_column`), so reserved words and names with spaces work, and `columns=[]` selects only the pseudo-columns instead of generating invalid SQL ([#40618](https://github.com/apache/beam/issues/40618)).

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,18 @@ 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")
+	// Remove this worker's directory on exit so that it does not accumulate on
+	// runners that reuse the environment for multiple pipelines. Only a direct
+	// child of semiPersistDir is ever removed.
+	if filepath.Dir(semiPersist) == filepath.Clean(*semiPersistDir) {
+		defer os.RemoveAll(semiPersist)
+	}
 	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 +268,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(), " "))