diff --git a/src/great_ai/great_ai/deploy/great_ai.py b/src/great_ai/great_ai/deploy/great_ai.py index c50caf0..462e7cd 100644 --- a/src/great_ai/great_ai/deploy/great_ai.py +++ b/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 ( - f"{self._version}+{get_function_metadata_store(self._func).model_versions}" - ) + flat_model_versions = ".".join(f"{k}-v{v}" for k, v in 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: diff --git a/src/great_ai/great_ai/models/load_model.py b/src/great_ai/great_ai/models/load_model.py deleted file mode 100644 index c8a133e..0000000 --- a/src/great_ai/great_ai/models/load_model.py +++ /dev/null @@ -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 diff --git a/src/great_ai/great_ai/models/use_model.py b/src/great_ai/great_ai/models/use_model.py index cd9ca3a..3e830ab 100644 --- a/src/great_ai/great_ai/models/use_model.py +++ b/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 diff --git a/src/great_ai/great_ai/views/function_metadata.py b/src/great_ai/great_ai/views/function_metadata.py index 97dc4b9..55369c0 100644 --- a/src/great_ai/great_ai/views/function_metadata.py +++ b/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