Add the directory with staged files to sys.path and document the usage

of `--files_to_stage` and `--beam_plugins`

Author: tvalentyn

.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
 }

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
 

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',

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')

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()

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:

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):
     """

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.