damccorm opened a new pull request, #40816:
URL: https://github.com/apache/beam/pull/40816
## Summary
Addresses #40673 using approach (b) from the issue discussion: spawn every
`MultiProcessShared` model in a dedicated process (`spawn_process=True`).
### Root Cause
When `MultiProcessShared(..., always_proxy=True, spawn_process=False)` is
used, the process that initializes the shared object both runs the
`_SingletonRegistrar` server (`Server.accepter`) on a background thread and
connects to its own server via `BaseProxy`. If an unreachable `BaseProxy` is
part of a reference cycle and Python's cyclic GC runs during an allocation on
the `accepter` thread, `BaseProxy._decref` opens a blocking client connection
to that same server while `accepter` is blocked inside `_decref`, deadlocking
the server (and any callers waiting on `_process_local_lock`).
### Changes
- **`apache_beam/ml/inference/base.py`**:
- Spawn shared models in dedicated processes (`spawn_process=True`) across
`_RunInferenceDoFn._load_model`, `_ModelHandlerManager.load`, and
`KeyedModelHandler.run_inference`, as well as shared `_ModelStatus` and
`_ModelRoutingStrategy` helpers (and use `always_proxy=False` for worker-local
`ModelManager` so `process_wide=True` metrics remain in the SDK worker without
self-proxying).
- Add `_TimedModelLoader` to propagate `ModelHandler` environment
variables, side-input model paths, `OOMProtectedFn` wrapping, and
`(load_model_latency_milli_secs, model_byte_size)` metrics from the spawned
model process back to the acquiring worker.
- Pass `destroy_on_zero_refcount=True` when `_ModelHandlerManager` evicts
per-key models (LRU eviction or `update_model_handler`) so evicted model server
processes exit cleanly once all in-flight inference references are released.
- **`apache_beam/utils/multi_process_shared.py`**:
- Serialize `constructor` with `apache_beam.internal.pickler` so lambdas,
closures, and `ModelHandler`s with local `inference_fn`s work with
`spawn_process=True`.
- Release entries directly via `singletonProxy_release_entry` on the proxy
connection (avoiding `RebuildProxy` inside the server process and avoiding
extra `AutoProxy` `incref`/`decref` sockets) and support
`destroy_on_zero_refcount` for graceful server shutdown once refcount reaches 0.
- Track and reap exited spawned processes and lifeline pipes
(`_reap_spawned_processes`), make `unsafe_hard_delete()` a safe no-op when a
tag does not exist, catch expected disconnect errors when a spawned server
exits via `os._exit(0)`, and invoke custom `__del__` finalizers before
`os._exit(0)` even if a timed-out call thread is still blocked in-flight.
- **Tests (`multi_process_shared_test.py` & `base_test.py`)**:
- Regression tests reproducing cyclic-GC-on-`accepter` with
reference-cycled `BaseProxy` handles and side-input model updates (#40673).
- Dedicated OS PID verification for single shared models, `model_copies >
1`, and multi-model `KeyedModelHandler`s, plus process termination verification
on LRU eviction, `update_model_handler`, and timeout hard-delete.
- Coverage for lambda/closure constructors, constructor error
cleanup/recovery, `destroy_on_zero_refcount` with active in-flight references,
`unsafe_hard_delete` during blocked calls, load latency/byte-size metrics, and
`env_vars` propagation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
