More experimentation

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

.vscode/settings.json

@@ -1,3 +1,19 @@
 {
-  "cSpell.words": ["pydantic", "pyplot", "sklearn", "Tfidf", "Vectorizer"]
+  "cSpell.words": [
+    "botocore",
+    "pydantic",
+    "pyplot",
+    "sklearn",
+    "Tfidf",
+    "tinydb",
+    "Vectorizer",
+    "xmargin"
+  ],
+  "files.exclude": {
+    ".env": true,
+    "**/.cache": true,
+    "**/.ipynb_checkpoints": true,
+    "**/.mypy_cache": true,
+    "**/.pytest_cache": true
+  }
 }

example/.gitignore

@@ -1,2 +1,3 @@
 data
 .ipynb_checkpoints
+tracing_database.json

example/README.md

@@ -0,0 +1,12 @@
+# Train Domain classifier from the [semantic scholar dataset](https://api.semanticscholar.org/corpus)
+
+## Upload the dataset (or a part of it) to shared infrastructure
+
+```sh
+mkdir ss-data && cd ss-data
+wget https://s3-us-west-2.amazonaws.com/ai2-s2-research-public/open-corpus/2022-02-01/manifest.txt
+wget -B https://s3-us-west-2.amazonaws.com/ai2-s2-research-public/open-corpus/2022-02-01/ -i manifest.txt
+cd -
+python3 -m good_ai.open_s3 --secrets s3.ini --push ss-data
+rm -rf ss-data
+```

example/main.py

@@ -1,15 +1,21 @@
 import json
+from random import shuffle
 
+from devtools import debug
 from predict_domain import predict_domain
 
-from good_ai import LargeFile, process_batch
+from good_ai import process_batch
 
 if __name__ == "__main__":
-    with open("data/s2-corpus-1583.json") as f:
+    with open(".cache/ss-data-0/s2-corpus-1583.json") as f:
         raw = json.load(f)
 
-    LargeFile.configure_credentials_from_file("s3.ini")
-
+    shuffle(raw)
     data = {f'{r["title"]} {r["abstract"]}': r["domain"] for r in raw[:5]}
 
-    print(process_batch(predict_domain, ["We have found a new type of chemical."]))
+    results = process_batch(predict_domain, data.keys())
+
+    for predicted, actual in zip(results, data.values()):
+        print(", ".join(actual))
+        debug(predicted)
+        print()

example/predict_domain.py

@@ -45,7 +45,7 @@ def predict_domain(
             )
         )
 
-        if sum(r.probability for r in results) >= cut_off_probability:
+        if sum(r.probability for r in results) >= cut_off_probability * 100:
             break
 
     return results

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