frank/more vat improvements (#194)

src/driftpy/market_map/market_map.py

@@ -103,7 +103,7 @@ def clear(self):
     def get_last_dump_filepath(self) -> str:
         return f"{market_type_to_string(self.market_type)}_{self.latest_slot}.pkl"
 
-    async def dump(self):
+    async def pre_dump(self) -> dict[str, bytes]:
         try:
             filters = []
             if is_variant(self.market_type, "Perp"):
@@ -142,13 +142,20 @@ async def dump(self):
                 raw_bytes = base64.b64decode(market["account"]["data"][0])
                 raw[str(pubkey)] = raw_bytes
 
+            return raw
+
+        except Exception as e:
+            print(f"error in marketmap pre-dump: {e}")
+
+    def dump(self, raw: dict[str, bytes], filename: Optional[str] = None):
+        try:
             markets = []
             for pubkey, market in raw.items():
                 markets.append(PickledData(pubkey=pubkey, data=compress(market)))
-            filename = (
+            path = filename or (
                 f"{market_type_to_string(self.market_type)}_{self.latest_slot}.pkl"
             )
-            with open(filename, "wb") as f:
+            with open(path, "wb") as f:
                 pickle.dump(markets, f)
 
         except Exception as e:

src/driftpy/pickle/vat.py

@@ -1,3 +1,4 @@
+import asyncio
 import pickle
 import os
 from dataclasses import dataclass
@@ -15,7 +16,7 @@ def __init__(
         self,
         drift_client: DriftClient,
         users: UserMap,
-        user_stats: Optional[UserStatsMap],
+        user_stats: UserStatsMap,
         spot_markets: MarketMap,
         perp_markets: MarketMap,
     ):
@@ -28,18 +29,37 @@ def __init__(
         self.perp_oracles = {}
         self.spot_oracles = {}
 
-    async def pickle(self):
-        await self.users.sync()
-        self.users.dump()
+    async def pickle(self, file_prefix: Optional[str] = None) -> dict[str, str]:
+        users_sync = asyncio.create_task(self.users.sync())
+        user_stats_sync = asyncio.create_task(self.user_stats.sync())
+        spot_markets_pre_dump = asyncio.create_task(self.spot_markets.pre_dump())
+        perp_markets_pre_dump = asyncio.create_task(self.perp_markets.pre_dump())
+        register_oracle_slot = asyncio.create_task(self.register_oracle_slot())
+
+        await asyncio.gather(
+            users_sync,
+            user_stats_sync,
+            spot_markets_pre_dump,
+            perp_markets_pre_dump,
+            register_oracle_slot,
+        )
+
+        spot_market_raw = spot_markets_pre_dump.result()
+        perp_market_raw = perp_markets_pre_dump.result()
+
+        filenames = self.get_filenames(file_prefix)
+
+        self.users.dump(filenames["users"])
 
-        if self.user_stats is not None:
-            await self.user_stats.sync()
-            self.user_stats.dump()
+        self.user_stats.dump(filenames["userstats"])
 
-        await self.spot_markets.dump()
-        await self.perp_markets.dump()
+        self.spot_markets.dump(spot_market_raw, filenames["spot_markets"])
 
-        await self.dump_oracles()
+        self.perp_markets.dump(perp_market_raw, filenames["perp_markets"])
+
+        self.dump_oracles(filenames["spot_oracles"], filenames["perp_oracles"])
+
+        return filenames
 
     async def unpickle(
         self,
@@ -56,8 +76,7 @@ async def unpickle(
         self.perp_markets.clear()
 
         await self.users.load(users_filename)
-        if self.user_stats is not None:
-            await self.user_stats.load(user_stats_filename)
+        await self.user_stats.load(user_stats_filename)
         await self.spot_markets.load(spot_markets_filename)
         await self.perp_markets.load(perp_markets_filename)
 
@@ -67,7 +86,12 @@ async def unpickle(
             self.spot_markets, self.perp_markets, self.spot_oracles, self.perp_oracles
         )
 
-    async def dump_oracles(self):
+    async def register_oracle_slot(self):
+        self.last_oracle_slot = (await self.drift_client.connection.get_slot()).value
+
+    def dump_oracles(
+        self, spot_filepath: Optional[str] = None, perp_filepath: Optional[str] = None
+    ):
         perp_oracles = []
         for market in self.drift_client.get_perp_market_accounts():
             oracle_price_data = self.drift_client.get_oracle_price_data_for_perp_market(
@@ -86,12 +110,12 @@ async def dump_oracles(self):
                 PickledData(pubkey=market.market_index, data=oracle_price_data)
             )
 
-        self.last_oracle_slot = (await self.drift_client.connection.get_slot()).value
-
-        with open(f"perporacles_{self.last_oracle_slot}.pkl", "wb") as f:
+        perp_path = perp_filepath or f"perporacles_{self.last_oracle_slot}.pkl"
+        with open(perp_path, "wb") as f:
             pickle.dump(perp_oracles, f)
 
-        with open(f"spotoracles_{self.last_oracle_slot}.pkl", "wb") as f:
+        spot_path = spot_filepath or f"spotoracles_{self.last_oracle_slot}.pkl"
+        with open(spot_path, "wb") as f:
             pickle.dump(spot_oracles, f)
 
     def load_oracles(
@@ -121,3 +145,29 @@ def load_oracles(
                     )  # oracle.pubkey is actually a market index
         else:
             raise FileNotFoundError(f"File {spot_filename} not found")
+
+    def get_filenames(self, prefix: Optional[str]) -> dict[str, str]:
+        filenames = {}
+
+        usermap_slot = self.users.get_slot()
+        userstats_slot = self.user_stats.latest_slot
+        spot_markets_slot = self.spot_markets.latest_slot
+        perp_markets_slot = self.perp_markets.latest_slot
+        oracle_slot = self.last_oracle_slot
+
+        if prefix:
+            filenames["users"] = f"{prefix}usermap_{usermap_slot}.pkl"
+            filenames["userstats"] = f"{prefix}userstats_{userstats_slot}.pkl"
+            filenames["spot_markets"] = f"{prefix}spot_{spot_markets_slot}.pkl"
+            filenames["perp_markets"] = f"{prefix}perp_{perp_markets_slot}.pkl"
+            filenames["spot_oracles"] = f"{prefix}spotoracles_{oracle_slot}.pkl"
+            filenames["perp_oracles"] = f"{prefix}perporacles_{oracle_slot}.pkl"
+        else:
+            filenames["users"] = f"usermap_{usermap_slot}.pkl"
+            filenames["userstats"] = f"userstats_{userstats_slot}.pkl"
+            filenames["spot_markets"] = f"spot_{spot_markets_slot}.pkl"
+            filenames["perp_markets"] = f"perp_{perp_markets_slot}.pkl"
+            filenames["spot_oracles"] = f"spotoracles_{oracle_slot}.pkl"
+            filenames["perp_oracles"] = f"perporacles_{oracle_slot}.pkl"
+
+        return filenames

src/driftpy/user_map/user_map.py

@@ -249,6 +249,6 @@ def dump(self, filename: Optional[str] = None):
         for pubkey, user in self.raw.items():
             users.append(PickledData(pubkey=pubkey, data=compress(user)))
         self.last_dumped_slot = self.get_slot()
-        filename = filename or f"usermap_{self.last_dumped_slot}.pkl"
-        with open(filename, "wb") as f:
+        path = filename or f"usermap_{self.last_dumped_slot}.pkl"
+        with open(path, "wb") as f:
             pickle.dump(users, f, pickle.HIGHEST_PROTOCOL)

src/driftpy/user_map/userstats_map.py

@@ -266,13 +266,13 @@ async def load(self, filename: Optional[str] = None):
                     Pubkey.from_string(str(user_stat.pubkey)), DataAndSlot(slot, data)
                 )
 
-    def dump(self):
+    def dump(self, filename: Optional[str] = None):
         user_stats = []
         for _pubkey, user_stat in self.raw.items():
             decoded: UserStatsAccount = decode_user_stat(user_stat)
             auth = decoded.authority
             user_stats.append(PickledData(pubkey=auth, data=compress(user_stat)))
         self.last_dumped_slot = self.latest_slot
-        filename = f"userstats_{self.last_dumped_slot}.pkl"
-        with open(filename, "wb") as f:
+        path = filename or f"userstats_{self.last_dumped_slot}.pkl"
+        with open(path, "wb") as f:
             pickle.dump(user_stats, f, pickle.HIGHEST_PROTOCOL)