Added basic error handling (#19)

Author: spietras

kilroy_module_pytorch_py_sdk/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kilroy-module-pytorch-py-sdk"
-version = "0.6.1"
+version = "0.6.2"
 description = "SDK for kilroy modules using PyTorch 🧰"
 readme = "README.md"
 authors = ["kilroy <[email protected]>"]

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/basic.py

@@ -1,3 +1,5 @@
+import json
+import logging
 from abc import ABC
 from dataclasses import dataclass
 from typing import (
@@ -43,6 +45,8 @@
     unpack_to_list,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class SupervisedLossMetric(Metric[Dict]):
     @classproperty
@@ -103,7 +107,6 @@ class State:
     generator: Generator
     codec: Codec
     results_cache: Dict[UUID, Tuple[Tensor, Tensor]]
-    used_results: Set[UUID]
     batch_size: int
     step: int
     metrics: MetricsState
@@ -181,12 +184,19 @@ async def generate(
         async for result in generated:
             sequences = unpack_to_list(result.sequences)
             for sequence, logprob in zip(sequences, result.logprobs):
+
                 post_id = uuid4()
+
                 async with self.state.read_lock() as state:
-                    post = await state.codec.encode(state.tokenizer, sequence)
+                    codec = state.codec
+                    tokenizer = state.tokenizer
+
+                post = await codec.encode(tokenizer, sequence)
+
                 if not dry:
                     async with self.state.write_lock() as state:
                         state.results_cache[post_id] = (sequence, logprob[0])
+
                 yield post_id, post
 
     async def _fit_supervised(self, data: AsyncIterable[Tensor]) -> None:
@@ -204,26 +214,38 @@ def fit(model, batch):
 
         async with batches.stream() as streamer:
             async for batch in streamer:
-                async with self.state.write_lock() as state:
-                    loss = await background(fit, state.model, batch)
-                    state.reports.step_supervised_losses.append(loss)
+                if batch:
+                    async with self.state.write_lock() as state:
+                        loss = await background(fit, state.model, batch)
+                        state.reports.step_supervised_losses.append(loss)
 
     async def fit_posts(
         self, posts: AsyncIterable[Tuple[Dict[str, Any], float]]
     ) -> None:
         async def decoded():
             async for post, _ in posts:
-                # noinspection PyShadowingNames
                 async with self.state.read_lock() as state:
-                    yield await state.codec.decode(state.tokenizer, post)
+                    codec = state.codec
+                    tokenizer = state.tokenizer
+                try:
+                    yield await codec.decode(tokenizer, post)
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to decode post: {json.dumps(post)}. Skipping...",
+                        exc_info=e,
+                    )
+                    continue
 
         await self._fit_supervised(decoded())
 
     async def _fit_reinforced(
         self,
         results: AsyncIterable[Tuple[Tensor, Tensor, Tensor]],
     ) -> None:
-        results = list([result async for result in results])
+        results = [result async for result in results]
+        if not results:
+            return
+
         logprobs = torch.stack([logprob for _, logprob, _ in results])
         scores = torch.stack([score for _, _, score in results])
 
@@ -239,10 +261,13 @@ def fit():
     async def fit_scores(self, scores: List[Tuple[UUID, float]]) -> None:
         async def get_results():
             for post_id, score in scores:
-                # noinspection PyShadowingNames
                 async with self.state.write_lock() as state:
+                    if post_id not in state.results_cache:
+                        logger.warning(
+                            f"Post {str(post_id)} has not been generated. Skipping..."
+                        )
+                        continue
                     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())
@@ -261,10 +286,8 @@ async def _reset_reports(state: State) -> None:
         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 _delete_results(state: State) -> None:
+        state.results_cache.clear()
 
     async def step(self) -> None:
         async with self.state.write_lock() as state:
@@ -284,5 +307,5 @@ async def step(self) -> None:
                 state.reports.step_reinforced_scores,
             )
             await self._reset_reports(state)
-            await self._delete_used_results(state)
+            await self._delete_results(state)
             state.step += 1

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/reward.py

@@ -1,3 +1,5 @@
+import json
+import logging
 from abc import ABC
 from asyncio import Queue, Task
 from dataclasses import dataclass
@@ -48,6 +50,8 @@
     unpack_to_list,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class SupervisedLossMetric(Metric[Dict]):
     @classproperty
@@ -165,7 +169,6 @@ 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
@@ -290,14 +293,19 @@ async def generate(
         async for result in generated:
             sequences = unpack_to_list(result.sequences)
             for sequence, logprob in zip(sequences, result.logprobs):
+
                 post_id = uuid4()
+
                 async with self.state.read_lock() as state:
-                    post = await state.codec.encode(
-                        state.language_model.tokenizer, sequence
-                    )
+                    codec = state.codec
+                    tokenizer = state.language_model.tokenizer
+
+                post = await codec.encode(tokenizer, sequence)
+
                 if not dry:
                     async with self.state.write_lock() as state:
                         state.results_cache[post_id] = (sequence, logprob[0])
+
                 yield post_id, post
 
     @staticmethod
@@ -350,6 +358,8 @@ async def _fit_supervised(
 
         async with batches.stream() as streamer:
             async for batch in streamer:
+                if not batch:
+                    continue
                 async with self.state.write_lock() as state:
                     sequences = pack_list(sequence for sequence, _ in batch)
                     scores = torch.vstack([score for _, score in batch])
@@ -372,13 +382,19 @@ async def fit_posts(
     ) -> None:
         async def decoded():
             async for post, score in posts:
-                # noinspection PyShadowingNames
                 async with self.state.read_lock() as state:
-                    post = await state.codec.decode(
-                        state.language_model.tokenizer, post
+                    codec = state.codec
+                    tokenizer = state.language_model.tokenizer
+                try:
+                    post = await codec.decode(tokenizer, post)
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to decode post: {json.dumps(post)}. Skipping...",
+                        exc_info=e,
                     )
-                    score = torch.tensor(score, dtype=torch.float)
-                    yield post, score
+                    continue
+                score = torch.tensor(score, dtype=torch.float)
+                yield post, score
 
         await self._fit_supervised(decoded())
 
@@ -421,6 +437,8 @@ async def _fit_reinforced(
 
         async with batches.stream() as streamer:
             async for batch in streamer:
+                if not batch:
+                    continue
                 sequences = pack_list([sequence for sequence, _, _ in batch])
                 scores = torch.vstack([score for _, _, score in batch])
                 async with self.state.write_lock() as state:
@@ -441,10 +459,13 @@ async def _fit_reinforced(
     async def fit_scores(self, scores: List[Tuple[UUID, float]]) -> None:
         async def get_results():
             for post_id, score in scores:
-                # noinspection PyShadowingNames
                 async with self.state.write_lock() as state:
+                    if post_id not in state.results_cache:
+                        logger.warning(
+                            f"Post {str(post_id)} has not been generated. Skipping..."
+                        )
+                        continue
                     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())
@@ -465,10 +486,8 @@ async def _reset_reports(state: State) -> None:
         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 _delete_results(state: State) -> None:
+        state.results_cache.clear()
 
     async def step(self) -> None:
         async with self.state.write_lock() as state:
@@ -503,5 +522,5 @@ async def step(self) -> None:
                 state.reports.step_reward_model_scores,
             )
             await self._reset_reports(state)
-            await self._delete_used_results(state)
+            await self._delete_results(state)
             state.step += 1

pyproject.toml

@@ -3,7 +3,7 @@
 
 [tool.poetry]
 name = "kilroy-module-pytorch-py-sdk"
-version = "0.6.1"
+version = "0.6.2"
 description = "SDK for kilroy modules using PyTorch 🧰"
 readme = "kilroy_module_pytorch_py_sdk/README.md"
 authors = ["kilroy <[email protected]>"]