Store model versions globally

src/great_ai/great_ai/deploy/great_ai.py

@@ -25,6 +25,7 @@ from ..helper import (
     snake_case_to_text,
     use_http_exceptions,
 )
+from ..models import model_versions
 from ..parameters import automatically_decorate_parameters
 from ..tracing.tracing_context import TracingContext
 from ..views import ApiMetadata, CacheStatistics, HealthCheckResponse, Trace
@@ -168,9 +169,11 @@ class GreatAI(Generic[T]):
 
     @property
     def version(self) -> str:
-        return (
+        flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions)
-            f"{self._version}+{get_function_metadata_store(self._func).model_versions}"
+        if flat_model_versions:
-        )
+            flat_model_versions = f"+{flat_model_versions}"
+
+        return f"{self._version}{flat_model_versions}"
 
     @property
     def documentation(self) -> str:

src/great_ai/great_ai/models/load_model.py

@@ -1,17 +0,0 @@
-from typing import Any, Optional, Tuple
-
-from dill import load
-
-from ..context import get_context
-
-
-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()
-
-    if path.is_dir():
-        return path, file.version
-
-    with file as f:
-        return load(f), file.version

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

src/great_ai/great_ai/views/function_metadata.py

@@ -6,5 +6,4 @@ from pydantic import BaseModel
 class FunctionMetadata(BaseModel):
     input_parameter_names: List[str] = []
     model_parameter_names: List[str] = []
-    model_versions: str = ""
     is_finalised: bool = False