Added support for saving state (#12)

Author: spietras

kilroy_module_pytorch_py_sdk/poetry.lock

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

kilroy_module_pytorch_py_sdk/pyproject.toml

@@ -1,11 +1,14 @@
+import json
 import random
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, AsyncIterable, Dict, Iterable, List, Set
 
 from kilroy_module_server_py_sdk import (
     CategorizableBasedParameter,
     Configurable,
     Parameter,
+    Savable,
     SerializableModel,
     classproperty,
 )
@@ -70,21 +73,52 @@ def parameters(cls) -> Set[Parameter]:
 
     async def build_default_state(self) -> State:
         params = Params(**self._kwargs)
-        sampler_cls = Sampler.for_category(params.sampler_type)
-        sampler_params = params.samplers_params.get(params.sampler_type, {})
-        if issubclass(sampler_cls, Configurable):
-            sampler = await sampler_cls.build(**sampler_params)
-            await sampler.init()
-        else:
-            sampler = sampler_cls(**sampler_params)
         return State(
-            sampler=sampler,
+            sampler=await self.build_generic(
+                Sampler,
+                category=params.sampler_type,
+                **params.samplers_params.get(params.sampler_type, {}),
+            ),
             samplers_params=params.samplers_params,
             contexts=params.contexts,
             max_length=params.max_length,
             batch_size=params.batch_size,
         )
 
+    async def save_state(self, state: State, directory: Path) -> None:
+        state_dict = {
+            "sampler_type": state.sampler.category,
+            "samplers_params": state.samplers_params,
+            "contexts": state.contexts,
+            "max_length": state.max_length,
+            "batch_size": state.batch_size,
+        }
+        if isinstance(state.sampler, Savable):
+            await state.sampler.save(directory / "sampler")
+        with open(directory / "state.json", "w") as f:
+            json.dump(state_dict, f)
+
+    async def load_saved_state(self, directory: Path) -> State:
+        with open(directory / "state.json", "r") as f:
+            state_dict = json.load(f)
+        sampler_type = state_dict["sampler_type"]
+        sampler_kwargs = {
+            **self._kwargs.get("samplers_params", {}).get(sampler_type, {}),
+            **state_dict["samplers_params"].get(sampler_type, {}),
+        }
+        return State(
+            sampler=await self.load_generic(
+                directory / "sampler",
+                Sampler,
+                category=state_dict["sampler_type"],
+                **sampler_kwargs,
+            ),
+            samplers_params=state_dict["samplers_params"],
+            contexts=state_dict["contexts"],
+            max_length=state_dict["max_length"],
+            batch_size=state_dict["batch_size"],
+        )
+
     async def cleanup(self) -> None:
         async with self.state.write_lock() as state:
             if isinstance(state.sampler, Configurable):

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/generator/generator.py

@@ -1,16 +1,7 @@
 from abc import ABC
 from asyncio import Queue, Task
 from dataclasses import dataclass
-from typing import (
-    Any,
-    AsyncIterable,
-    Coroutine,
-    Dict,
-    Generator,
-    List,
-    Set,
-    Tuple,
-)
+from typing import Any, AsyncIterable, Coroutine, Dict, List, Set, Tuple
 from uuid import UUID, uuid4
 
 import numpy as np
@@ -32,6 +23,7 @@
 from torch.nn import MSELoss, NLLLoss
 from torch.nn.utils.rnn import PackedSequence
 
+from kilroy_module_pytorch_py_sdk import Generator
 from kilroy_module_pytorch_py_sdk.codec import Codec
 from kilroy_module_pytorch_py_sdk.models import LanguageModel, RewardModel
 from kilroy_module_pytorch_py_sdk.optimizers import Optimizer

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/reward.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Dict, List
 
+import torch
 from kilroy_module_server_py_sdk import (
     Configurable,
     SerializableModel,
@@ -66,8 +68,21 @@ def schema(cls) -> Dict[str, Any]:
             return {"type": "number", "minimum": 0}
 
     async def build_default_state(self) -> State:
+        model_params = self._kwargs.pop("parameters")
         user_params = Params(**self._kwargs)
-        return State(optimizer=Adam(self._params, **user_params.dict()))
+        return State(optimizer=Adam(model_params, **user_params.dict()))
+
+    @classmethod
+    async def save_state(cls, state: State, directory: Path) -> None:
+        with open(directory / "optimizer.pt", "wb") as f:
+            await background(torch.save, state.optimizer.state_dict(), f)
+
+    async def load_saved_state(self, directory: Path) -> State:
+        with open(directory / "optimizer.pt", "rb") as f:
+            state_dict = await background(torch.load, f)
+        optimizer = Adam(self._kwargs.pop("parameters"))
+        optimizer.load_state_dict(state_dict)
+        return State(optimizer=optimizer)
 
     async def step(self) -> None:
         async with self.state.write_lock() as state:

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/optimizers/adam.py

@@ -54,12 +54,6 @@ async def undo() -> None:
 
 
 class Optimizer(Categorizable, ABC):
-    def __init__(
-        self, params: Union[Iterable[Tensor], Iterable[Dict]], *args, **kwargs
-    ) -> None:
-        super().__init__(*args, **kwargs)
-        self._params = list(params)
-
     @classproperty
     def category(cls) -> str:
         name: str = cls.__name__

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/optimizers/base.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Dict
 
+import torch
 from kilroy_module_server_py_sdk import (
     Configurable,
     SerializableModel,
@@ -55,8 +57,21 @@ def schema(cls) -> Dict[str, Any]:
             return {"type": "number", "minimum": 0}
 
     async def build_default_state(self) -> State:
+        model_params = self._kwargs.pop("parameters")
         user_params = Params(**self._kwargs)
-        return State(optimizer=RMSprop(self._params, **user_params.dict()))
+        return State(optimizer=RMSprop(model_params, **user_params.dict()))
+
+    @classmethod
+    async def save_state(cls, state: State, directory: Path) -> None:
+        with open(directory / "optimizer.pt", "wb") as f:
+            await background(torch.save, state.optimizer.state_dict(), f)
+
+    async def load_saved_state(self, directory: Path) -> State:
+        with open(directory / "optimizer.pt", "rb") as f:
+            state_dict = await background(torch.load, f)
+        optimizer = RMSprop(self._kwargs.pop("parameters"))
+        optimizer.load_state_dict(state_dict)
+        return State(optimizer=optimizer)
 
     async def step(self) -> None:
         async with self.state.write_lock() as state:

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/optimizers/rmsprop.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Dict
 
+import torch
 from kilroy_module_server_py_sdk import (
     Configurable,
     SerializableModel,
@@ -49,8 +51,21 @@ def schema(cls) -> Dict[str, Any]:
             return {"type": "number", "minimum": 0}
 
     async def build_default_state(self) -> State:
+        model_params = self._kwargs.pop("parameters")
         user_params = Params(**self._kwargs)
-        return State(optimizer=SGD(self._params, **user_params.dict()))
+        return State(optimizer=SGD(model_params, **user_params.dict()))
+
+    @classmethod
+    async def save_state(cls, state: State, directory: Path) -> None:
+        with open(directory / "optimizer.pt", "wb") as f:
+            await background(torch.save, state.optimizer.state_dict(), f)
+
+    async def load_saved_state(self, directory: Path) -> State:
+        with open(directory / "optimizer.pt", "rb") as f:
+            state_dict = await background(torch.load, f)
+        optimizer = SGD(self._kwargs.pop("parameters"))
+        optimizer.load_state_dict(state_dict)
+        return State(optimizer=optimizer)
 
     async def step(self) -> None:
         async with self.state.write_lock() as state:

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/optimizers/sgd.py

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