Repository navigation
[fix] Keep inherited Ray lock handles alive in forked children - #202
OutstanderWang wants to merge 1 commit into
Conversation
…leasing them Releasing an inherited Ray actor handle in a fork child runs the parent's core worker there. The child crashes in _int_free, the parent's lock actors can be reported out of scope and die, and the parent's later GCS calls time out. A datasets map with num_proc after tq.init() hit all three: every TransferQueueLockManager died and ray.get_actor in the parent timed out. _reset_after_fork now parks the inherited _owned, _managers and _leases in _inherited_after_fork, which is never released, before resetting the state. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 170f6f1eec
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| _managers: dict = {} # shard -> TransferQueueLockManager handle, looked up on first use | ||
| _owned: list = [] # the lock actors this process's tq.init() created, killed by its tq.close() | ||
| _closes = 0 # this process's tq.close() calls: a lease created before the latest one is given up | ||
| _inherited_after_fork: list = [] # a forked child's copies of the parent's handles, never released |
There was a problem hiding this comment.
This trailing comment is a one-line description of what the variable contains and how long it lives, not the reason that retaining the references is necessary; the scoped comment rule requires a concise 2–4-line “why” explanation. _reset_after_fork already provides that rationale immediately before the append, so delete this duplicate rather than leave two explanations that can drift.
AGENTS.md reference: AGENTS.md:L24-L26
Useful? React with 👍 / 👎.
OutstanderWang
changed the title
170f6f1 to
fbd4236
Compare
| fake_lock_actors.options.assert_not_called() | ||
|
|
||
|
|
||
| def test_fork_child_keeps_inherited_lock_handles_referenced(monkeypatch): |
There was a problem hiding this comment.
Would it be helpful to add a test for handles retained through _leases as well?
2 participants
What this PR does
After
tq.init(), a process that forks (for exampledatasets.map(num_proc=N)) can break both the child and the parent when it holds handles to the KV global lock actors.The child inherits the parent's Ray actor handles for the lock:
_owned,_managers, and the handles held by the leases in_leases._reset_after_fork()used to just reassign these globals to empty containers. That dropped the last reference to the inherited handles, so they were released inside the child. Releasing a Ray handle there runs the parent's core worker in the forked process, which causes three problems:_int_free.TransferQueueLockManagerdied.ray.get_actor(...).When it triggers
Only a process that holds lock-actor handles when it forks is affected:
tq.init()started TransferQueue withTQ_ENABLE_KV_GLOBAL_LOCKset, since_ownedholds the actors it created; orkv_global_lock, which looks the actors up and caches them in_managers. If it forks while holding a lock,_leasescarries handles too.Without the global lock,
tq.init()creates no lock actors and_owned,_managersand_leasesstay empty, so this bug cannot trigger.kv_local_lockis in-process only and holds no Ray handles.Fix
Before resetting its state,
_reset_after_fork()now moves the inherited_owned,_managersand_leasesinto a module-level_inherited_after_forklist, and nothing ever releases that list. The child still starts with a clean lock state, but the inherited handles keep a reference and are never garbage-collected in the child, so Ray never tries to release them there.The cost is that a forked child holds on to a few stale handle objects for its whole lifetime. That is negligible. Within TransferQueue it is the only option, because an inherited handle cannot be released safely from a fork. Callers can avoid the issue entirely by not forking after
tq.init(), for example by using thespawnstart method.This relies on the child not running module teardown. A child started with
multiprocessing'sforkmethod exits throughos._exit, so the list is never collected. A child that instead exits through normal interpreter shutdown could still release the handles during teardown. That is Ray's general limitation on forking afterray.init(), and it is out of scope here.Changes
transfer_queue/kv_lock.py: adds_inherited_after_forkand moves the inherited state into it in_reset_after_fork().tests/test_kv_lock.py: addstest_fork_child_keeps_inherited_lock_handles_referenced. It runs_reset_after_fork()andgc.collect(), then checks two things: the live state is reset, and the objects inherited through_ownedand_managersare still alive (checked throughweakref)._leasesis parked the same way but is not covered by the test.How to reproduce (before this fix)
TQ_ENABLE_KV_GLOBAL_LOCKfor thetq.init()that starts TransferQueue. Then fork either from that process, or from any process that has already calledkv_global_lock.datasets.Dataset.map(..., num_proc>1).TransferQueueLockManageractors die, and a laterray.get_actorin the parent times out.Test