Store model versions globally

src/great_ai/great_ai/models/use_model.py

@@ -1,11 +1,25 @@
 from functools import wraps
-from typing import Any, Callable, Dict, List, Literal, TypeVar, Union, cast
+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
-from .load_model import load_model
 
 F = TypeVar("F", bound=Callable[..., Any])
 
@@ -13,14 +27,14 @@ F = TypeVar("F", bound=Callable[..., Any])
 def use_model(
     key: str,
     *,
-    version: Union[int, Literal["latest"]],
+    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 version"
+    ), "Only integers or the string literal `latest` is allowed as a version"
 
-    model, actual_version = load_model(
+    model, actual_version = _load_model(
         key=key,
         version=None if version == "latest" else version,
     )
@@ -30,9 +44,6 @@ def use_model(
 
         store = get_function_metadata_store(func)
         store.model_parameter_names.append(model_kwarg_name)
-        if store.model_versions:
-            store.model_versions += "."
-        store.model_versions += f"{key}-v{actual_version}"
 
         @wraps(func)
         def wrapper(*args: List[Any], **kwargs: Dict[str, Any]) -> Any:
@@ -44,3 +55,19 @@ def use_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