More experimentation

Signed-off-by: András Schmelczer <andras@schmelczer.dev>

good_ai/src/good_ai/good_ai/__init__.py

@@ -1,2 +1,3 @@
 from .deploy import process_batch
 from .models import save_model, use_model
+from .set_default_config import set_default_config, set_default_config_if_uninitialized

good_ai/src/good_ai/good_ai/core/__init__.py

@@ -1,2 +0,0 @@
-from .function_registry import function_registry
-from .plugin import Plugin

good_ai/src/good_ai/good_ai/core/function_registry.py

@@ -1,23 +0,0 @@
-from collections import defaultdict
-from typing import Any, Callable, DefaultDict, List
-
-from .plugin import Plugin
-
-
-class FunctionRegistry:
-    def __init__(self) -> None:
-        self._registered_functions: DefaultDict[int, List[Plugin]] = defaultdict(
-            lambda: []
-        )
-
-    def add_plugin(self, function: Callable[..., Any], plugin: Plugin):
-        self._registered_functions[id(function)].append(plugin)
-
-    def get_plugins(self, function: Callable[..., Any]) -> List[Plugin]:
-        plugins = self._registered_functions[id(function)]
-        for p in plugins:
-            p.initialize()
-        return plugins
-
-
-function_registry = FunctionRegistry()

good_ai/src/good_ai/good_ai/core/plugin.py

@@ -1,19 +0,0 @@
-from typing import Any, Callable
-
-
-class Plugin:
-    def __init__(self, function: Callable[[Any], Any]):
-        self._function = function
-        self._initialized = False
-
-    def initialize(self):
-        if not self._initialized:
-            self.on_initialize()
-            self._initialized = True
-
-    def on_initialize(self):
-        pass
-
-    def __call__(self, *args: Any, **kwargs: Any) -> Any:
-        assert self._initialized
-        return self._function(*args, **kwargs)

good_ai/src/good_ai/good_ai/deploy/process_batch.py

@@ -1,16 +1,23 @@
-from functools import partial, reduce
 from typing import Any, Callable, Iterable, Optional, Sequence
 
 from good_ai.utilities.parallel_map import parallel_map
 
-from ..core import function_registry
+from ..set_default_config import set_default_config_if_uninitialized
+from ..tracing import Trace, TracingContext
 
 
 def process_batch(
     function: Callable[..., Any],
     batch: Iterable[Any],
     concurrency: Optional[int] = None,
-) -> Sequence[Any]:
-    plugins = function_registry.get_plugins(function)
-    composed = partial(reduce, lambda r, f: f(r), plugins)
-    return parallel_map(composed, batch, concurrency=concurrency)
+) -> Sequence[Trace]:
+    set_default_config_if_uninitialized()
+
+    def inner(input: Any) -> Trace:
+        with TracingContext() as t:
+            t.log_input(input)
+            result = function(input)
+            output = t.log_output(result)
+        return output
+
+    return parallel_map(inner, batch, concurrency=concurrency)

good_ai/src/good_ai/good_ai/models/load_model.py

@@ -1,20 +1,24 @@
 import logging
-from typing import Any, Optional
+from typing import Any, Optional, Tuple
 
 from joblib import load
 
 from good_ai.open_s3 import LargeFile
 
+from ..set_default_config import set_default_config_if_uninitialized
+
 logger = logging.getLogger("models")
 
 
 def load_model(
     key: str, version: Optional[int] = None, return_path: bool = False
-) -> Any:
+) -> Tuple[Any, int]:
+    set_default_config_if_uninitialized()
+
     file = LargeFile(name=key, mode="rb", version=version)
 
     if return_path:
-        return file.get()
+        return file.get(), file.version
 
     with file as f:
-        return load(f)
+        return load(f), file.version

good_ai/src/good_ai/good_ai/models/save_model.py

@@ -6,12 +6,16 @@ from joblib import dump
 
 from good_ai.open_s3 import LargeFile
 
+from ..set_default_config import set_default_config_if_uninitialized
+
 logger = logging.getLogger("models")
 
 
 def save_model(
     model: Union[Path, str, object], key: str, keep_last_n: Optional[int] = None
 ) -> int:
+    set_default_config_if_uninitialized()
+
     file = LargeFile(name=key, mode="wb", keep_last_n=keep_last_n)
 
     if isinstance(model, Path) or isinstance(model, str):

good_ai/src/good_ai/good_ai/models/use_model.py

@@ -1,24 +1,33 @@
-from typing import Literal, Union
+from functools import wraps
+from typing import Any, Callable, Literal, Union
 
-from ..core import function_registry
-from .use_model_plugin import UseModelPlugin
+from ..tracing import Model, TracingContext
+from .load_model import load_model
 
 
 def use_model(
-    key: str, version: Union[int, Literal["latest"]] = None, return_path: bool = False
-):
+    key: str,
+    *,
+    version: Union[int, Literal["latest"]],
+    return_path: bool = False,
+    model_kwarg_name: str = "model"
+) -> Callable[..., Any]:
     assert isinstance(version, int) or version == "latest"
 
-    def inner(f):
-        function_registry.add_plugin(
-            f,
-            UseModelPlugin(
-                f,
-                key=key,
-                version=version if isinstance(version, int) else None,
-                return_path=return_path,
-            ),
-        )
-        return f
+    model, actual_version = load_model(
+        key=key,
+        version=None if version == "latest" else version,
+        return_path=return_path,
+    )
 
-    return inner
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            context = TracingContext.get_current_context()
+            if context:
+                context.log_model(Model(key=key, version=actual_version))
+            return func(*args, **kwargs, **{model_kwarg_name: model})
+
+        return wrapper
+
+    return decorator

good_ai/src/good_ai/good_ai/models/use_model_plugin.py

@@ -1,28 +0,0 @@
-from typing import Any, Callable, Optional
-
-from ..core import Plugin
-from .load_model import load_model
-
-
-class UseModelPlugin(Plugin):
-    def __init__(
-        self,
-        function: Callable[[Any], Any],
-        key: str,
-        version: Optional[int],
-        return_path: bool,
-    ):
-        def wrapper(*args, **kwargs):
-            return function(*args, **kwargs, model=self._model)
-
-        super().__init__(wrapper)
-
-        self._key = key
-        self._version = version
-        self._return_path = return_path
-
-    def on_initialize(self) -> None:
-        super().on_initialize()
-        self._model = load_model(
-            key=self._key, version=self._version, return_path=self._return_path
-        )

good_ai/src/good_ai/good_ai/set_default_config.py

@@ -0,0 +1,55 @@
+import logging
+import random
+from pathlib import Path
+
+from good_ai.good_ai.tracing.tracing_context import TracingContext
+from good_ai.open_s3 import LargeFile
+
+from .tracing import PersistenceDriver, TinyDbDriver
+
+logger = logging.getLogger("good_ai")
+
+_initialized = False
+
+
+def set_default_config_if_uninitialized() -> None:
+    if not _initialized:
+        set_default_config()
+
+
+def set_default_config(
+    log_level: int = logging.INFO,
+    s3_config: Path = Path("s3.ini"),
+    seed: int = 42,
+    tracing_db_driver: PersistenceDriver = TinyDbDriver(Path("tracing_database.json")),
+) -> None:
+    global _initialized
+    logging.basicConfig(level=log_level)
+
+    _initialize_large_file(s3_config)
+    _set_seed(seed)
+
+    TracingContext.persistence_driver = tracing_db_driver
+
+    _initialized = True
+
+    logger.info(f"Defaults: configured ✅")
+
+
+def _initialize_large_file(s3_config: Path) -> None:
+    if s3_config.exists():
+        LargeFile.configure_credentials_from_file(s3_config)
+    else:
+        logger.info(
+            f"Provided S3 config ({s3_config.resolve()}) not found, skipping LargeFile initialisation"
+        )
+
+
+def _set_seed(seed: int) -> None:
+    random.seed(seed)
+    try:
+        import numpy
+
+        numpy.random.seed(seed + 1)
+    except ImportError:
+        pass

good_ai/src/good_ai/good_ai/tracing/__init__.py

@@ -0,0 +1,4 @@
+from .model import Model
+from .persistence import MongoDbDriver, PersistenceDriver, TinyDbDriver
+from .trace import Trace
+from .tracing_context import TracingContext

good_ai/src/good_ai/good_ai/tracing/model.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class Model(BaseModel):
+    key: str
+    version: int

good_ai/src/good_ai/good_ai/tracing/persistence/__init__.py

@@ -0,0 +1,3 @@
+from .mongodb_driver import MongoDbDriver
+from .persistence_driver import PersistenceDriver
+from .tinydb_driver import TinyDbDriver

good_ai/src/good_ai/good_ai/tracing/persistence/mongodb_driver.py

@@ -0,0 +1,5 @@
+from .persistence_driver import PersistenceDriver
+
+
+class MongoDbDriver(PersistenceDriver):
+    pass

good_ai/src/good_ai/good_ai/tracing/persistence/persistence_driver.py

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+
+class PersistenceDriver(ABC):
+    @abstractmethod
+    def save_document(self, document: Dict[str, Any]) -> str:
+        pass

good_ai/src/good_ai/good_ai/tracing/persistence/tinydb_driver.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+from typing import Any, Dict
+
+from tinydb import TinyDB
+
+from .persistence_driver import PersistenceDriver
+
+
+class TinyDbDriver(PersistenceDriver):
+    def __init__(self, path_to_db: Path) -> None:
+        super().__init__()
+        self._db = TinyDB(path_to_db)
+
+    def save_document(self, document: Dict[str, Any]) -> str:
+        return self._db.insert(document)

good_ai/src/good_ai/good_ai/tracing/trace.py

@@ -0,0 +1,14 @@
+from typing import Any, List
+
+from pydantic import BaseModel
+
+from .model import Model
+
+
+class Trace(BaseModel):
+    created: str
+    execution_time_ms: float
+    input: Any
+    models: List[Model]
+    output: Any
+    evaluation: Any = None

good_ai/src/good_ai/good_ai/tracing/tracing_context.py

@@ -0,0 +1,71 @@
+import logging
+import threading
+from collections import defaultdict
+from datetime import datetime
+from types import TracebackType
+from typing import Any, DefaultDict, List, Optional, Type
+
+from .persistence import PersistenceDriver
+
+logger = logging.getLogger("good_ai")
+
+from .model import Model
+from .trace import Trace
+
+
+class TracingContext:
+    persistence_driver: PersistenceDriver
+    _contexts: DefaultDict[int, List["TracingContext"]] = defaultdict(lambda: [])
+
+    def __init__(self) -> None:
+        self._models: List[Model] = []
+        self._input: Any = None
+        self._output: Any = None
+        self._trace: Optional[Trace] = None
+        self._start_time = datetime.utcnow()
+
+    def log_input(self, input: Any) -> None:
+        self._input = input
+
+    def log_model(self, model: Model) -> None:
+        self._models.append(model)
+
+    def log_output(self, output: Any) -> Trace:
+        self._output = output
+
+        delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
+        self._trace = Trace(
+            created=self._start_time.isoformat(),
+            execution_time_ms=delta_time,
+            input=self._input,
+            models=self._models,
+            output=self._output,
+        )
+        return self._trace
+
+    @classmethod
+    def get_current_context(cls) -> Optional["TracingContext"]:
+        if cls._contexts[threading.get_ident()]:
+            return cls._contexts[threading.get_ident()][-1]
+        return None
+
+    def __enter__(self) -> "TracingContext":
+        self._contexts[threading.get_ident()].append(self)
+        return self
+
+    def __exit__(
+        self,
+        type: Optional[Type[BaseException]],
+        exception: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        assert self._contexts[threading.get_ident()][-1] == self
+        self._contexts[threading.get_ident()].remove(self)
+
+        if type is None:
+            assert self._trace is not None
+            self.persistence_driver.save_document(self._trace.dict())
+        else:
+            logger.exception(f"Could not finish operation: {exception}")
+
+        return True

good_ai/src/good_ai/open_s3/__main__.py

@@ -3,8 +3,8 @@
 import logging
 from pathlib import Path
 
-from large_file import LargeFile
-from parse_arguments import parse_arguments
+from .large_file import LargeFile
+from .parse_arguments import parse_arguments
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)