Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 30 additions & 1 deletion python/tvm/contrib/popen_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -366,19 +366,48 @@ def __init__(
self._maximum_process_uses = maximum_process_uses
self._stdout = stdout
self._stderr = stderr
self._shutdown = False

if self._initializer is not None and not callable(self._initializer):
raise TypeError("initializer must be callable for PopenPoolExecutor")

def __del__(self):
"""Destructor.

Note
----
Called during garbage collection. This may be called later than expected.
Always call shutdown() explicitly to avoid deadlocks.
"""
if not self._shutdown:
self.shutdown(wait=True)

def shutdown(self, wait=True):
"""Shutdown the executor and clean up resources.

Parameters
----------
wait : bool
If True, wait for pending work to complete.

Note
----
DEADLOCK WARNING: This method can deadlock when called during garbage
collection due to exception reference cycles. When exceptions occur,
Python creates reference cycles that delay garbage collection. The
deadlock happens when: exception creates reference cycle → new pool
creates worker → GC cleans old pool → old pool's __del__ calls shutdown()
which tries to acquire locks again.
"""
self._lock.acquire()
for worker in self._worker_map.values():
try:
worker.kill()
except ImportError:
pass
self._lock.release()
self._threadpool.shutdown()
self._threadpool.shutdown(wait=wait)
self._shutdown = True

def _worker_run(self, fn, args, kwargs):
"""Internal thread runner."""
Expand Down
4 changes: 2 additions & 2 deletions python/tvm/meta_schedule/builder/local_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,7 +192,7 @@ def build(self, build_inputs: List[BuilderInput]) -> List[BuilderResult]:
)
else:
raise ValueError("Unreachable: unexpected result: {map_result}")
del pool
pool.shutdown()
return results

def _sanity_check(self) -> None:
Expand All @@ -208,7 +208,7 @@ def _check(f_build, f_export) -> None:
)
value = pool.submit(_check, self.f_build, self.f_export)
value.result()
del pool
pool.shutdown()


def _worker_func(
Expand Down
Loading