Skip to content

Support closures and clean hard-deletion in MultiProcessShared spawn_process (#40673) - #40822

Open
damccorm wants to merge 1 commit into
apache:masterfrom
damccorm:mps-spawn-process-fixes
Open

damccorm wants to merge 1 commit into
apache:masterfrom
damccorm:mps-spawn-process-fixes

Conversation

@damccorm

@damccorm damccorm commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Prerequisite MultiProcessShared(..., spawn_process=True) improvements for #40673, split out ahead of the RunInference changes:

  1. Serialize constructors with apache_beam.internal.pickler: multiprocessing.get_context('spawn') uses standard-library pickle for Process(args=...), which rejects lambdas, closures, and local callables (AttributeError: Can't get local object ...). Serializing self._constructor with apache_beam.internal.pickler allows MultiProcessShared(lambda: ..., spawn_process=True) and objects holding local functions to spawn cleanly.
  2. Make MultiProcessShared.unsafe_hard_delete() safe for spawned processes:
    • Return immediately as a no-op if the tag is not registered locally and has no .address file, avoiding accidentally spawning a new server process just to delete it.
    • Ignore expected (EOFError, OSError, BrokenPipeError) disconnects when the spawned server exits via os._exit(0) inside handle_unsafe_hard_delete.
  3. Run custom __del__ finalizers before os._exit(0) during in-flight calls: If an in-flight RPC thread (such as a timed-out call) still holds a stack reference to self.obj when _SingletonEntry.unsafe_hard_delete() runs in a spawned server process, del self.obj alone does not drop the refcount to zero before os._exit(0) terminates the process. Explicitly invoking type(obj).__del__(obj) ensures custom cleanup hooks still execute.

Testing

  • Added unit tests in MultiProcessSharedSpawnProcessTest (multi_process_shared_test.py) covering:
    • Dedicated OS PID isolation across tags and shared PID reuse for matching tags
    • Cyclic-GC-on-Server.accepter deadlock regression ([Bug]: Using MultiProcessShared handle can deadlock when always_proxy=True #40673) with always_proxy=True, spawn_process=True
    • Lambda and closure constructors with spawn_process=True
    • unsafe_hard_delete() on non-existent tags (no-op, no extra process spawned)
    • unsafe_hard_delete() running __del__ while a blocked RPC is in flight

@damccorm
damccorm force-pushed the mps-spawn-process-fixes branch from 367bec6 to 5d02372 Compare October 26, 2026 19:00
@damccorm
damccorm marked this pull request as ready for review October 26, 2026 19:01

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📋 Review Summary

This PR prepares MultiProcessShared(..., spawn_process=True) for broader usage (issue #40673) by serializing constructors with apache_beam.internal.pickler, preventing unsafe_hard_delete() from spawning new processes on nonexistent tags, and ensuring custom __del__ finalizers run before os._exit(0). The implementation is clean and well-tested; however, the sys.getrefcount(obj) > 2 check in _SingletonEntry.unsafe_hard_delete is brittle on Python 3.14+ where LOAD_FAST_BORROW optimizes local variable references.

🔍 General Feedback

  • Python 3.14 Compatibility (sys.getrefcount): Starting in Python 3.14 (PEP 709 / LOAD_FAST_BORROW bytecode optimization), passing a local variable obj to sys.getrefcount(obj) uses a borrowed reference on the evaluation stack rather than incrementing the refcount. When no in-flight RPC holds a reference, sys.getrefcount(obj) returns 1 instead of 2 (or 2 instead of 3 when one _SingletonProxy or in-flight RPC reference exists). Calling type(obj).__del__(obj) unconditionally (or moving del obj after a weakref.ref check) avoids relying on CPython bytecode refcount semantics.

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

# still hold a stack reference to `obj`, preventing CPython from
# invoking `obj.__del__()` before `os._exit(0)`.
if (self._hard_delete_callback is not None and
sys.getrefcount(obj) > 2 and hasattr(type(obj), '__del__')):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium: sys.getrefcount(obj) > 2 relies on CPython incrementing the reference count both for the local variable obj and for passing obj as an argument to sys.getrefcount(). In Python 3.14+, the compiler emits LOAD_FAST_BORROW for local variable loads, which does not increment the reference count when pushing obj onto the stack for sys.getrefcount(obj). As a result:

  • When only one in-flight thread holds a reference to obj (total actual references = 2: obj local + the in-flight frame), sys.getrefcount(obj) returns 2 on Python 3.14+, so > 2 evaluates to False and __del__ will be skipped before os._exit(0).
  • Conversely, if del obj fails to drop the refcount to zero (and os._exit(0) is about to terminate the process anyway when self._hard_delete_callback is not None), we can avoid sys.getrefcount altogether by checking a weakref after del obj or invoking type(obj).__del__(obj) before self._hard_delete_callback() exits the process.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Potentially check for other holders with a weakref instead (drop our reference, then see whether ref() still returns the object)?

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @claudevdm for label python.
Backup reviewer: jrmccluskey (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

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

if (self._hard_delete_callback is not None and
sys.getrefcount(obj) > 2 and hasattr(type(obj), '__del__')):
try:
type(obj).__del__(obj)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this hang for a model handler with a lock on the model + with_exception_handling(timeout ?

  1. The worker calls predict.
  2. It runs in the model process, takes the model's lock and hangs.
  3. The timeout fires in the worker.
  4. The worker stops waiting and schedules the model for deletion. Predict is still hung and still holding the lock.
  5. The worker loads a fresh copy of the model and carries on.
  6. At least 2× the timeout later, finish_bundle sends "hard delete" to the old model process. Forget the model, notice predict still holds it, and call the model's cleanup first. The cleanup waits for the lock, which predict never releases, so os._exit never runs and finish_bundle waits forever for a reply.

Potentially run the cleanup on a separate thread and wait for it only up to a timeout, then exit regardless, so a cleanup stuck behind the hung call can't block the hard delete?

if (self._hard_delete_callback is not None and
sys.getrefcount(obj) > 2 and hasattr(type(obj), '__del__')):
try:
type(obj).__del__(obj)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The model's cleanup potentially runs twice with with_exception_handling(timeout) if the hung call finishes while the cleanup is still running?

  1. A call on the model hangs. The timeout fires, and the worker schedules the model for deletion. The call is still running in the model process.
  2. Later, the hard delete sees the call still holds the model (sys.getrefcount(obj) > 2), so it calls type(obj).__del__(obj) (run 1).
  3. The cleanup makes the hung call end, e.g. by killing a subprocess the call was polling. Once the call ends and its worker thread exits, nothing else holds the model.
  4. Run 1 returns and del obj drops the last reference. Python frees the model and runs __del__ again (run 2), because it doesn't know about run 1.

A cleanup that isn't safe to repeat would then run twice. Potentially keep the model referenced until exit after running the cleanup (e.g. append it to a module-level list), so __del__ can't run a second time?

# still hold a stack reference to `obj`, preventing CPython from
# invoking `obj.__del__()` before `os._exit(0)`.
if (self._hard_delete_callback is not None and
sys.getrefcount(obj) > 2 and hasattr(type(obj), '__del__')):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Potentially check for other holders with a weakref instead (drop our reference, then see whether ref() still returns the object)?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants