From 8044b4935d8bd7baef2a1270aeb4585d9fb52aac Mon Sep 17 00:00:00 2001 From: Rohin Bhasin Date: Mon, 15 Apr 2024 19:46:57 -0400 Subject: [PATCH] Change `clear_local` to be parallellized. --- runhouse/servers/obj_store.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/runhouse/servers/obj_store.py b/runhouse/servers/obj_store.py index f0cb81ef2..77d7013d1 100644 --- a/runhouse/servers/obj_store.py +++ b/runhouse/servers/obj_store.py @@ -931,9 +931,10 @@ async def aclear_for_env_servlet_name(env_servlet_name: str): async def aclear_local(self): if self.has_local_storage: - for k in list(self._kv_store.keys()): - # Pop handles removing from global obj store vs local one - await self.apop_local(k) + # Use asyncio gather to run all the deletes concurrently + await asyncio.gather( + *[self.apop_local(k) for k in list(self._kv_store.keys())] + ) async def aclear(self): logger.warning("Clearing all keys from all envs in the object store!")