Move files

great_ai/models/__init__.py

@@ -0,0 +1,2 @@
+from .save_model import save_model
+from .use_model import model_versions, use_model

great_ai/models/save_model.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+from typing import Optional, Union
+
+from dill import dump
+
+from ..context import get_context
+
+
+def save_model(
+    model: Union[Path, str, object], key: str, *, keep_last_n: Optional[int] = None
+) -> str:
+    file = get_context().large_file_implementation(
+        name=key, mode="wb", keep_last_n=keep_last_n
+    )
+
+    if isinstance(model, Path) or isinstance(model, str):
+        file.push(model)
+    else:
+        with file as f:
+            dump(model, f)
+
+    get_context().logger.info(f"Model {key} uploaded with version {file.version}")
+
+    return f"{key}:{file.version}"

great_ai/models/use_model.py

@@ -0,0 +1,73 @@
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Set,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from dill import load
+
+from ..context import get_context
+from ..helper import get_function_metadata_store
+from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
+from ..tracing.tracing_context import TracingContext
+from ..views import Model
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def use_model(
+    key: str,
+    *,
+    version: Union[int, Literal["latest"]] = "latest",
+    model_kwarg_name: str = "model",
+) -> Callable[[F], F]:
+    assert (
+        isinstance(version, int) or version == "latest"
+    ), "Only integers or the string literal `latest` is allowed as a version"
+
+    model, actual_version = _load_model(
+        key=key,
+        version=None if version == "latest" else version,
+    )
+
+    def decorator(func: F) -> F:
+        assert_function_is_not_finalised(func)
+
+        store = get_function_metadata_store(func)
+        store.model_parameter_names.append(model_kwarg_name)
+
+        @wraps(func)
+        def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
+            tracing_context = TracingContext.get_current_tracing_context()
+            if tracing_context:
+                tracing_context.log_model(Model(key=key, version=actual_version))
+            return func(*args, **kwargs, **{model_kwarg_name: model})
+
+        return cast(F, wrapper)
+
+    return decorator
+
+
+model_versions: Set[Tuple[str, int]] = set()
+
+
+def _load_model(key: str, version: Optional[int] = None) -> Tuple[Any, int]:
+    file = get_context().large_file_implementation(name=key, mode="rb", version=version)
+    path = file.get()
+
+    model_versions.add((key, file.version))
+
+    if path.is_dir():
+        return path, file.version
+
+    with file as f:
+        return load(f), file.version