Refactor _SharedCache to handle context vs non-context ownership

[shunping]

Currently, _SharedCache registers all live owners as joint owners of every cached subprocess key. This works fine when all owners are within a unified context (e.g. a test class wrapping multiple runs), but it causes unintended resource sharing/leakage when there are concurrent, independent, non-context owners (such as a long-lived runner).

Here is a resource leak scenario:

1. The Prism runner starts and registers a persistent non-context owner owner_1.
2. The pipeline requests an external transform, spinning up a short-lived Java Expansion Service that registers its own owner owner_2 and requests its startup command key.
3. In the old implementation, get() automatically added all currently registered live owners to the cache key's owner set.
   

   beam/sdks/python/apache_beam/utils/subprocess_server.py

   Lines 117 to 118 in 361c2c4

   	for owner in self._live_owners:
   	self._cache[key].owners.add(owner)

   
   As a result, owner_1 (the Prism runner) was registered as a joint owner of the short-lived Expansion Service's cache key.
4. When the pipeline finished using the external transform and called purge(owner_2), the short-lived Expansion Service subprocess was never terminated because owner_1 was still registered as an owner (entry.owners is not empty).
   

   beam/sdks/python/apache_beam/utils/subprocess_server.py

   Lines 104 to 105 in 361c2c4

   	if not entry.owners:
   	to_delete.append(entry.obj)

5. This resulted in orphaned, zombie Expansion Service processes accumulating in the background, binding up localhost ports and exhausting resources.

This PR refactors _SharedCache in subprocess_server.py to explicitly distinguish between context owners (e.g., cache_subprocesses which should hold ownership over everything created under its block) and non-context owners (e.g., individual subprocess server instances like PrismRunner or expansion service subprocesses).

Additionally, get() now takes an optional owner arg. Context owners automatically own all created keys, while independent non-context owners only own the keys they explicitly request.

---

For example, we run YAML example test with pytest apache_beam/yaml/examples/testing/examples_test.py --log-cli-level=INFO -s.

• Without this PR, we see all expansion service servers created in these 46 tests and prism server being destroyed at the end of pytest execution.

  Really destroying service at localhost:49623 with cmd: ('/Users/shunping/.apache_beam/cache/prism/bin/prism', '--job_port', '49623', '--log_level', 'info', '--log_kind', 'json', '--serve_http=false')
  WARNING:apache_beam.utils.subprocess_server:Really destroying service at localhost:49745 with cmd: ['/Users/shunping/.apache_beam/cache/venvs/e08b2501f55abf367196f6578e3dbed1fe2a6d9244fc96ad5d82f3d5c621c1cf/bin/python', '-m', 'apache_beam.runners.portability.expansion_service_main', '--port', '49745', '--fully_qualified_name_glob=*', '--pickle_library=cloudpickle', '--requirements_file=/Users/shunping/.apache_beam/cache/venvs/e08b2501f55abf367196f6578e3dbed1fe2a6d9244fc96ad5d82f3d5c621c1cf-requirements.txt', '--serve_loopback_worker']
  WARNING:apache_beam.utils.subprocess_server:Really destroying service at localhost:49855 with cmd: ['/Users/shunping/.apache_beam/cache/venvs/5052f7801853faa5017d022495ad2a30f1ae0a321fa50de56227e75028f0e612/bin/python', '-m', 'apache_beam.runners.portability.expansion_service_main', '--port', '49855', '--fully_qualified_name_glob=*', '--pickle_library=cloudpickle', '--requirements_file=/Users/shunping/.apache_beam/cache/venvs/5052f7801853faa5017d022495ad2a30f1ae0a321fa50de56227e75028f0e612-requirements.txt', '--serve_loopback_worker']
  WARNING:apache_beam.utils.subprocess_server:Really destroying service at localhost:49888 with cmd: ['/Users/shunping/.apache_beam/cache/venvs/aedc40ba446114c3eb019ef9c411c0ab2204da67a2f4c732530c1c6aba7eef31/bin/python', '-m', 'apache_beam.runners.portability.expansion_service_main', '--port', '49888', '--fully_qualified_name_glob=*', '--pickle_library=cloudpickle', '--requirements_file=/Users/shunping/.apache_beam/cache/venvs/aedc40ba446114c3eb019ef9c411c0ab2204da67a2f4c732530c1c6aba7eef31-requirements.txt', '--serve_loopback_worker']
  

• With this PR, each expansion service server is destroyed as soon as the corresponding test is over. Prism server is cleaned up before pytest ends.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a resource leak issue in the subprocess management logic where independent subprocess owners were incorrectly sharing ownership of cached keys. By introducing a distinction between context-based and independent owners, the system now ensures that resources are correctly tracked and cleaned up, preventing zombie processes and port exhaustion.

Highlights

• Refactored _SharedCache ownership model: Updated _SharedCache to distinguish between context and non-context owners, preventing unintended resource sharing and leakage.
• Improved resource lifecycle management: Subprocess keys are now only owned by the requesting non-context owner, ensuring that short-lived services are properly terminated upon purge.
• Updated API and testing: Modified the register and get methods to support the new ownership logic and added comprehensive unit tests to verify the behavior.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request updates the _SharedCache implementation in subprocess_server.py to transition _live_owners from a set to a dictionary, allowing the system to distinguish between context and non-context owners. It also updates the get and register methods to handle this ownership model and adds corresponding unit tests. The review feedback correctly identifies a critical race condition and a potential resource leak in the get method, where the _live_owners check is performed outside the lock and the cache entry is created before validating the requesting owner. Implementing the suggested fix to perform these checks inside the lock is highly recommended.

[shunping]

assign set of reviewers

[github-actions]

Assigning reviewers:

R: @jrmccluskey for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).