Delayed results cleaning to step (#18)

Author: spietras

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/basic.py

@@ -103,6 +103,7 @@ class State:
     generator: Generator
     codec: Codec
     results_cache: Dict[UUID, Tuple[Tensor, Tensor]]
+    used_results: Set[UUID]
     batch_size: int
     step: int
     metrics: MetricsState
@@ -240,7 +241,8 @@ async def get_results():
             for post_id, score in scores:
                 # noinspection PyShadowingNames
                 async with self.state.write_lock() as state:
-                    sequence, logprob = state.results_cache.pop(post_id)
+                    sequence, logprob = state.results_cache.get(post_id)
+                    state.used_results.add(post_id)
                 yield sequence, logprob, torch.tensor(score)
 
         await self._fit_reinforced(get_results())
@@ -258,6 +260,12 @@ async def _reset_reports(state: State) -> None:
         state.reports.step_supervised_losses = []
         state.reports.step_reinforced_scores = []
 
+    @staticmethod
+    async def _delete_used_results(state: State) -> None:
+        for post_id in state.used_results:
+            state.results_cache.pop(post_id, None)
+        state.used_results.clear()
+
     async def step(self) -> None:
         async with self.state.write_lock() as state:
             await state.optimizer.step()
@@ -276,4 +284,5 @@ async def step(self) -> None:
                 state.reports.step_reinforced_scores,
             )
             await self._reset_reports(state)
+            await self._delete_used_results(state)
             state.step += 1

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/reward.py

@@ -165,6 +165,7 @@ class State:
     backend_generator: Generator
     codec: Codec
     results_cache: Dict[UUID, Tuple[Tensor, Tensor]]
+    used_results: Set[UUID]
     batch_size: int
     sample_size: int
     step: int
@@ -442,7 +443,8 @@ async def get_results():
             for post_id, score in scores:
                 # noinspection PyShadowingNames
                 async with self.state.write_lock() as state:
-                    sequence, logprob = state.results_cache.pop(post_id)
+                    sequence, logprob = state.results_cache.get(post_id)
+                    state.used_results.add(post_id)
                 yield sequence, logprob, torch.tensor(score)
 
         await self._fit_reinforced(get_results())
@@ -462,6 +464,12 @@ async def _reset_reports(state: State) -> None:
         state.reports.step_reward_model_losses = []
         state.reports.step_reward_model_scores = []
 
+    @staticmethod
+    async def _delete_used_results(state: State) -> None:
+        for post_id in state.used_results:
+            state.results_cache.pop(post_id, None)
+        state.used_results.clear()
+
     async def step(self) -> None:
         async with self.state.write_lock() as state:
             await state.language_model.optimizer.step()
@@ -495,4 +503,5 @@ async def step(self) -> None:
                 state.reports.step_reward_model_scores,
             )
             await self._reset_reports(state)
+            await self._delete_used_results(state)
             state.step += 1