FIX Callback transport don't store connections forever - #34870
FIX Callback transport don't store connections forever#34870jeremiedbb wants to merge 4 commits into
Conversation
| # Nothing more will be sent for this fit. The channel itself outlives the fit, | ||
| # since this callback keeps logging across fits, but the connection a worker | ||
| # copy holds is released here rather than whenever that copy is collected. | ||
| self._channel.disconnect() |
There was a problem hiding this comment.
That part is new and useful to stop the leak because callbacks are held a reference cycles and thus are not collected immediately. Adding the fact that loky reuses worker makes the callbacks and thus the channels and thus the connections alive for a very long time, which gives time for the number of file descriptors to grow too much.
So ideally the better fix would be to avoid that reference cycle (it comes from the context tree holding references to both its parent and children), which I'm planning to investigate as a follow-up work. In the meantime the disconnect fix is acceptable.
| def test_channel_state_is_not_walked_by_the_pickler(factory, monkeypatch): | ||
| """Check that pickling a callback never traverses state its channels write to. | ||
|
|
||
| An estimator carrying a callback can be pickled by a background thread, e.g. loky's | ||
| queue feeder dispatching a task to a worker, while the listener thread of that same | ||
| queue feeder dispatching a task to a worker, while a channel thread of that same | ||
| callback mutates the callback's state as messages come in. Pickling a container that | ||
| another thread mutates breaks the dump, which joblib reports as "Could not pickle | ||
| the task to send it to the workers". | ||
|
|
||
| The check runs from a hook, i.e. while the listeners are up, which is when such a | ||
| The check runs from a hook, i.e. while the channels are up, which is when such a | ||
| dispatch would happen. |
There was a problem hiding this comment.
I had to rewrite this test because it used to inspect the global _message_consumers dict which is no longer there. Instead I monkeypatch Channel.init to record the message_consumer. The intent of the test is still the same and it still fails if we remove the fix from #34821
| - Fixed a leak of file descriptors and background threads by | ||
| :class:`callback.ScoringMonitor` and :class:`callback.ProgressBar`. | ||
| By :user:`Jérémie du Boisberranger <jeremiedbb>`. :pr:`34856`, |
There was a problem hiding this comment.
this looks weird but the idea is to reference this PR and #34856 in a single changelog entry. Look at the generated changelog, it looks as expected.
|
The missing coverage is for a line that was already uncovered before the refactoring (see https://app.codecov.io/gh/scikit-learn/scikit-learn/pull/34870). I'd rather ignore it for this PR and open a follow up PR to write a test to cover this line (unrelated to this PR). |
Fixes #34734
This PR fixes the second (and main) source of leaking threads and file descriptors. It ensures that we don't keep open connections forever, even when they're not useful anymore.
(I'd rather have #34856 merged first to have easier merge conflicts but the 2 PRs fix 2 independent problems and can be reviewed independently)
Currently, we store every connection in a process-wide dict (
_worker_connections) and never remove them. So they stay open until the worker is shut down (which can take a while since loky can reuse workers for subsequent tasks). Making it a lru cache mitigates the issue a bit, but not entirely since it still allows to grow the number of open connections far beyond what's actually needed.This PR proposes to refactor the transport module a bit to add a new object (Channel) that abstracts both end of the communication channel (listener in the main process and clients in the workers). This way the connection made to
sendmessages is tied to the channel and dies with it. The connection alone can also be closed to immediately free the resources instead of waiting for the callback to be collected (i.e. at teardown for scoring monitor).Maybe it could have been fixed with a bit less refactoring but I actually find it cleaner and simpler. Now the only public object from transport is
Channel(at least when we make transport public :) ). The callbacks create and store a channel and then call methods of the channel.