TL;DR: itertools.tee keeps an item referenced until its whole internal batch (57 items currently) is consumed, even after every branch has consumed the individual item. Therefore, __del__ and any resource release (e.g. GPU memory for CUDA tensors) happen later than expected. Two questions: (1) is a short doc note about this welcome, and (2) is there interest in changing this behavior, to free items as soon as all branches pass?
Note: This is a follow-up based on issue #138765 and PR #154369, which were closed. I’d like to reopen the discussion here rather than discussing on the closed items, I hope that’s alright. Happy to move to a different category if that’s more appropriate.
The issue was filed as a “memory leak”, which is not the case IMO. I believe this is about finalization timing. An item which is already consumed by all branches stays referenced until its whole batch is retired. So Py_DECREF and __del__ are delayed, in a non-trivial and surprising way. The original reporter’s issue was PyTorch CUDA tensors: the tensor’s GPU memory stays pinned long after every iterator has passed the item.
Two questions
1. Docs: Currently the docs state:
This itertool may require significant auxiliary storage (depending on how much temporary data needs to be stored). In general, if one iterator uses most or all of the data before another iterator starts, it is faster to use
list()instead oftee().
This note is about storage size, it says nothing about delayed finalization of already-consumed items, and limits the note to the data that “needs to be stored”. I’d read this as the gap between the slowest and fastest iterator.
Would a short, implementation-independent sentence or change be welcome, e.g. something like: “Items already consumed by every iterator may not be released immediately, their finalizers can be delayed until internal buffers are recycled”? I understand the reluctance to over-specify (I filed PR #154369 pinning the 57 figure, which was rejected as an implementation detail).
2. Behavior: In the issue @storchaka describes a fix that releases items as soon as every branch has passed them, without adding memory or time cost.
My questions:
- Is there interest in this behavioral change, or is the current design considered final?
- @storchaka do you intend to implement the proposal yourself? If not, I’m happy to give it a shot.
- If the design is final, is there some way to make this behavior more clear in the docs?
Thanks for your time and considerations!
cc @rhettinger
PS: A small example showing this in action:
from itertools import tee
class LoudLifecycle:
def __init__(self, i):
self.i = i
print(f"INIT {self.i}")
def __del__(self):
print(f"DEL {self.i}")
def my_iter():
for i in range(100):
yield LoudLifecycle(i)
a, b = tee(my_iter())
for _ in zip(a, b):
pass