Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes Dendrite new loop close #2654

Merged
Merged
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
18 changes: 14 additions & 4 deletions bittensor/core/dendrite.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,14 +163,18 @@ async def session(self) -> aiohttp.ClientSession:
self._session = aiohttp.ClientSession()
return self._session

def close_session(self):
def close_session(self, using_new_loop: bool = False):
"""
Closes the internal `aiohttp <https://github.com/aio-libs/aiohttp>`_ client session synchronously.

This method ensures the proper closure and cleanup of the aiohttp client session, releasing any
resources like open connections and internal buffers. It is crucial for preventing resource leakage
and should be called when the dendrite instance is no longer in use, especially in synchronous contexts.

Arguments:
using_new_loop: A flag to determine whether this has been called with a new event loop rather than
the default. This will indicate whether to close this event loop at the end of this call.

Note:
This method utilizes asyncio's event loop to close the session asynchronously from a synchronous context.
It is advisable to use this method only when asynchronous context management is not feasible.
Expand All @@ -188,6 +192,8 @@ def close_session(self):
if self._session:
loop = asyncio.get_event_loop()
loop.run_until_complete(self._session.close())
if using_new_loop:
loop.close()
self._session = None

async def aclose_session(self):
Expand Down Expand Up @@ -374,16 +380,20 @@ def query(
category=DeprecationWarning,
)
result = None
use_new_loop = False
try:
loop = asyncio.get_event_loop()
result = loop.run_until_complete(self.forward(*args, **kwargs))
except Exception:
except Exception as e:
logging.debug(
f"Exception encountered while running Dendrite.query initially: {e}"
)
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
result = new_loop.run_until_complete(self.forward(*args, **kwargs))
new_loop.close()
use_new_loop = True
finally:
self.close_session()
self.close_session(using_new_loop=use_new_loop)
return result # type: ignore

async def forward(
Expand Down