diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index 00c0cdc3f9c8..89cec619b020 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": "38701", - "modification": 55 + "modification": 56 } diff --git a/CHANGES.md b/CHANGES.md index 59c7ac7b24ba..1e027c997479 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -70,7 +70,7 @@ * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). * (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)). -* X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* (Python) Staged files directory is now automatically added to `sys.path` on the Python SDK worker at startup. This makes Python files provided via the '--files_to_stage' pipeline option importable in the pipeline code and makes it easier to initialize Python SDK harness at startup via the `--beam_plugins` pipeline option. For more information, see the [Staging Individual Files](https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/#staging-files) section of the dependency management docs. This behavior can be disabled by passing the '--experiments=no_staged_dir_in_sys_path' pipeline option ([#39431](https://github.com/apache/beam/issues/39431)). ## Breaking Changes diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index afbec3f46f71..239d577cfa0e 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1863,6 +1863,7 @@ def _add_argparse_args(cls, parser): 'workers will install them in same order they were specified on ' 'the command line.')) parser.add_argument( + '--file_to_stage', '--files_to_stage', dest='files_to_stage', action='append', diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py index 90ad27d0a39b..b321314b4b76 100644 --- a/sdks/python/apache_beam/options/pipeline_options_test.py +++ b/sdks/python/apache_beam/options/pipeline_options_test.py @@ -569,6 +569,22 @@ def test_extra_package(self): options = PipelineOptions(flags=['']) self.assertEqual(options.get_all_options()['extra_packages'], None) + def test_files_to_stage(self): + options = PipelineOptions([ + '--file_to_stage', + 'abc', + '--files_to_stage', + 'def', + '--files_to_stage', + 'ghi' + ]) + self.assertEqual( + sorted(options.get_all_options()['files_to_stage']), + ['abc', 'def', 'ghi']) + + options = PipelineOptions(flags=['']) + self.assertEqual(options.get_all_options()['files_to_stage'], None) + def test_dataflow_job_file(self): options = PipelineOptions(['--dataflow_job_file', 'abc']) self.assertEqual(options.get_all_options()['dataflow_job_file'], 'abc') diff --git a/sdks/python/apache_beam/runners/portability/beam_plugins_it_test.py b/sdks/python/apache_beam/runners/portability/beam_plugins_it_test.py new file mode 100644 index 000000000000..54984503fa54 --- /dev/null +++ b/sdks/python/apache_beam/runners/portability/beam_plugins_it_test.py @@ -0,0 +1,70 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import logging +import os +import shutil +import tempfile +import unittest +import uuid + +import pytest + +import apache_beam as beam +from apache_beam.options.pipeline_options import SetupOptions +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to + +_LOGGER = logging.getLogger(__name__) + + +class BeamPluginsIT(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + + def tearDown(self): + shutil.rmtree(self.temp_dir) + + @pytest.mark.it_postcommit + def test_beam_plugins_staging(self): + pipeline = TestPipeline(is_integration_test=True) + setup_options = pipeline.options.view_as(SetupOptions) + + plugin_name = f'beam_integration_plugin_{uuid.uuid4().hex[:8]}' + plugin_file_path = os.path.join(self.temp_dir, f'{plugin_name}.py') + + with open(plugin_file_path, 'w') as f: + f.write("import sys\nsys.beam_plugin_loaded_for_test = True\n") + + staged_files = setup_options.files_to_stage or [] + staged_files.append(plugin_file_path) + setup_options.files_to_stage = staged_files + setup_options.beam_plugins = [plugin_name] + + def check_plugin_loaded(_): + import sys + return getattr(sys, 'beam_plugin_loaded_for_test', False) + + with pipeline as p: + res = (p | beam.Create([None]) | beam.Map(check_plugin_loaded)) + assert_that(res, equal_to([True])) + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + unittest.main() diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py index 754a631eaf33..1c464141f7e3 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py @@ -50,6 +50,7 @@ _LOGGER = logging.getLogger(__name__) _ENABLE_GOOGLE_CLOUD_PROFILER = 'enable_google_cloud_profiler' _FN_LOG_HANDLER = None +_STAGED_DIRECTORY = 'staged' def _import_beam_plugins(plugins): @@ -131,6 +132,12 @@ def create_harness(environment, dry_run=False): environment.get('RUNNER_CAPABILITIES', '').split()) _LOGGER.info('semi_persistent_directory: %s', semi_persistent_directory) + experiments = sdk_pipeline_options.view_as(DebugOptions).experiments or [] + if 'no_staged_dir_in_sys_path' not in experiments and semi_persistent_directory: + staged_dir = os.path.join(semi_persistent_directory, _STAGED_DIRECTORY) + if os.path.isdir(staged_dir) and staged_dir not in sys.path: + sys.path.append(staged_dir) + _worker_id = environment.get('WORKER_ID', None) try: diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py b/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py index 5ecd9616fcf9..c84438953b13 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_main_test.py @@ -129,6 +129,55 @@ def __init__(self, *unused): def test_import_beam_plugins(self): sdk_worker_main._import_beam_plugins(BeamPlugin.get_all_plugin_paths()) + def test_create_harness_adds_staged_dir_to_sys_path(self): + import sys + import tempfile + + with tempfile.TemporaryDirectory() as temp_dir: + staged_dir = os.path.join(temp_dir, sdk_worker_main._STAGED_DIRECTORY) + os.mkdir(staged_dir) + + env = { + 'CONTROL_API_SERVICE_DESCRIPTOR': '', + 'SEMI_PERSISTENT_DIRECTORY': temp_dir, + } + + if staged_dir in sys.path: + sys.path.remove(staged_dir) + + sdk_worker_main.create_harness(env, dry_run=True) + + try: + self.assertIn(staged_dir, sys.path) + finally: + if staged_dir in sys.path: + sys.path.remove(staged_dir) + + def test_create_harness_does_not_add_staged_dir_with_experiment(self): + import sys + import tempfile + + with tempfile.TemporaryDirectory() as temp_dir: + staged_dir = os.path.join(temp_dir, sdk_worker_main._STAGED_DIRECTORY) + os.mkdir(staged_dir) + + env = { + 'CONTROL_API_SERVICE_DESCRIPTOR': '', + 'SEMI_PERSISTENT_DIRECTORY': temp_dir, + 'PIPELINE_OPTIONS': '{"experiments":["no_staged_dir_in_sys_path"]}', + } + + if staged_dir in sys.path: + sys.path.remove(staged_dir) + + sdk_worker_main.create_harness(env, dry_run=True) + + try: + self.assertNotIn(staged_dir, sys.path) + finally: + if staged_dir in sys.path: + sys.path.remove(staged_dir) + @staticmethod def _overrides_case_to_option_dict(case): """ diff --git a/website/www/site/content/en/documentation/sdks/python-pipeline-dependencies.md b/website/www/site/content/en/documentation/sdks/python-pipeline-dependencies.md index 26dc1db1bc06..7d27f91d322d 100644 --- a/website/www/site/content/en/documentation/sdks/python-pipeline-dependencies.md +++ b/website/www/site/content/en/documentation/sdks/python-pipeline-dependencies.md @@ -91,9 +91,48 @@ If your pipeline uses packages that are not available publicly (e.g. packages th See the [build documentation](https://pypa-build.readthedocs.io/en/latest/index.html) for more details on this command. +## Staging Individual Files {#staging-files} + +If your pipeline relies on one or more individual Python files or non-Python data files that do not need to be packaged as a full Python package, you can stage them individually to the remote workers. + +To stage individual files, run your pipeline with the `--files_to_stage` (or `--file_to_stage`) pipeline option. This option accepts a list of local file paths: + + --files_to_stage="/path/to/my_module.py,/path/to/data_config.json" + +When the pipeline runs, the runner uploads these files and makes them available on the workers in the worker's staged files directory. + +### Accessing Staged Files on the Workers + +Staged files are downloaded onto the worker: + +* **Python modules**: Starting with Apache Beam 2.76.0, to make it easy to import staged Python files as modules, Beam automatically appends the worker's staged files directory to the Python interpreter's search path (`sys.path`) during startup. This means you can import them directly in your code: + + import my_module + +* **Data or configuration files**: If you staged non-Python files (such as a JSON config), they are downloaded to the staged files directory. You can locate this directory on the worker by reading the `SEMI_PERSISTENT_DIRECTORY` environment variable, or look for `/tmp/staged` which is the default location for staged files on containerized runners: + + import os + staged_dir = os.environ.get('SEMI_PERSISTENT_DIRECTORY', '/tmp/staged') + config_path = os.path.join(staged_dir, 'data_config.json') + +### Importing Plugins on Worker Startup + +You can use staged files in combination with the `--beam_plugins` pipeline option (supported starting with Apache Beam 2.76.0) to run initialization code on the workers before any processing starts. + +To use a staged file as a plugin: +1. Stage the plugin file (e.g. `my_custom_plugin.py`) using `--file_to_stage`. +2. Reference the module name in `--beam_plugins`. + +For example, run your pipeline with: + + --file_to_stage="/path/to/my_custom_plugin.py" \ + --beam_plugins="my_custom_plugin" + +This instructs the SDK worker to import `my_custom_plugin` immediately on startup, triggering any initialization logic defined in the module. + ## Multiple File Dependencies {#multiple-file-dependencies} -Often, your pipeline code spans multiple files. To run your project remotely, you must group these files as a Python package and specify the package when you run your pipeline. When the remote workers start, they will install your package. To group your files as a Python package and make it available remotely, perform the following steps: +Often, your pipeline code spans multiple files. To run your project remotely, it is best to group these files as a Python package and specify the package when you run your pipeline. When the remote workers start, they will install your package. To group your files as a Python package and make it available remotely, perform the following steps: 1. Create a [setup.py](https://pythonhosted.org/an_example_pypi_project/setuptools.html) file for your project. The following is a very basic `setup.py` file.