From 30f6d8e7eefdbbb2fb9c09afad7e84a200bf9097 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Tue, 20 Oct 2026 19:40:56 -0700 Subject: [PATCH 1/5] Support Environment Sensitive Staging --- .../runners/dataflow/internal/apiclient.py | 32 +++--- .../dataflow/internal/apiclient_test.py | 108 +++++++++++++++++- sdks/python/container/boot.go | 11 +- 3 files changed, 129 insertions(+), 22 deletions(-) diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py index 8a478f3fb633..cc28a65e792e 100644 --- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient.py +++ b/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 - 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(): diff --git a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py b/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py index b82c7f550673..6d1c34612b7b 100644 --- a/sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py +++ b/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) diff --git a/sdks/python/container/boot.go b/sdks/python/container/boot.go index 1e912ebebb75..409bf2b67bba 100644 --- a/sdks/python/container/boot.go +++ b/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 /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(), " ")) From 73038aca44f12edc638712751a1378146a11fc55 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Wed, 21 Oct 2026 11:23:43 -0700 Subject: [PATCH 2/5] Add Change commit --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index d9ca6f0b83dd..cf102ef3fd86 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -116,6 +116,7 @@ * (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that directly follows an escaped delimiter, which merged two records into one ([#40459](https://github.com/apache/beam/issues/40459)). * (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) 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)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes From 40447359a986677f6d73cf99e33e2a74ded2ca12 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Fri, 23 Oct 2026 13:59:18 -0700 Subject: [PATCH 3/5] trigger post submits --- .github/trigger_files/beam_PostCommit_Python.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index b6d790320d0b..4ecf816cf37e 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.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": 65 + "pr": "40699", + "modification": 66 } From 4bec8d5e055db2fceaa82c55e9dc6017d62cda5e Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Fri, 23 Oct 2026 14:41:05 -0700 Subject: [PATCH 4/5] trivial change --- .github/trigger_files/beam_PostCommit_Python.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index 4ecf816cf37e..edf20146ce9a 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.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": "40699", - "modification": 66 + "modification": 67 } From 1870a428550ca2fb8fa1ca1df66d8663e61fb405 Mon Sep 17 00:00:00 2001 From: Tarun Annapareddy Date: Sun, 25 Oct 2026 15:22:57 -0700 Subject: [PATCH 5/5] add clean up and update changes --- .github/trigger_files/beam_PostCommit_Python.json | 2 +- CHANGES.md | 2 ++ sdks/python/container/boot.go | 6 ++++++ 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index edf20146ce9a..35d8c450e24a 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.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": "40699", - "modification": 67 + "modification": 68 } diff --git a/CHANGES.md b/CHANGES.md index 4c70a013fe8c..fb5130b2e6f0 100644 --- a/CHANGES.md +++ b/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 `//staged` instead of `/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 diff --git a/sdks/python/container/boot.go b/sdks/python/container/boot.go index 409bf2b67bba..dbc108ed9396 100644 --- a/sdks/python/container/boot.go +++ b/sdks/python/container/boot.go @@ -217,6 +217,12 @@ func launchSDKProcess() error { // 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)