Skip to content

Reference#

from great_ai import *

Core#

GreatAI #

Bases: Generic[T, V]

Source code in great_ai/deploy/great_ai.py
class GreatAI(Generic[T, V]):
    __name__: str
    __doc__: str

    def __init__(
        self,
        func: Callable[..., Union[V, Awaitable[V]]],
    ):
        """Do not call this function directly, use GreatAI.create instead."""

        func = automatically_decorate_parameters(func)
        get_function_metadata_store(func).is_finalised = True

        self._cached_func = self._get_cached_traced_function(func)
        self._wrapped_func = wraps(func)(freeze_arguments(self._cached_func))

        wraps(func)(self)
        self.__doc__ = (
            f"GreatAI wrapper for interacting with the `{self.__name__}` "
            + f"function.\n\n{dedent(self.__doc__ or '')}"
        )

        self.version = str(get_context().version)
        flat_model_versions = ".".join(f"{k}-v{v}" for k, v in model_versions)
        if flat_model_versions:
            self.version += f"+{flat_model_versions}"

        self.app = FastAPI(
            title=snake_case_to_text(self.__name__),
            version=self.version,
            description=self.__doc__
            + f"\n\nFind out more in the [dashboard]({DASHBOARD_PATH}).",
            docs_url=None,
            redoc_url=None,
        )

        self._bootstrap_rest_api()

    @overload
    @staticmethod
    def create(  # type: ignore
        # Overloaded function signatures 1 and 2 overlap with incompatible return types
        # https://github.com/python/mypy/issues/12759
        func: Callable[..., Awaitable[V]],
    ) -> "GreatAI[Awaitable[Trace[V]], V]":
        ...

    @overload
    @staticmethod
    def create(
        func: Callable[..., V],
    ) -> "GreatAI[Trace[V], V]":
        ...

    @staticmethod
    def create(
        func: Union[Callable[..., Awaitable[V]], Callable[..., V]],
    ) -> Union["GreatAI[Awaitable[Trace[V]], V]", "GreatAI[Trace[V], V]"]:
        """Decorate a function by wrapping it in a GreatAI instance.

        The function can be typed, synchronous or async. If it has
        unwrapped parameters (parameters not affected by a @parameter
        or @use_model decorator), those will be automatically wrapped.

        The return value is replaced by a Trace (or Awaitable[Trace]),
        while the original return value is available under the `.output`
        property.

        For configuration options, see great_ai.configure.

        Examples:
            >>> @GreatAI.create
            ... def my_function(a):
            ...     return a + 2
            >>> my_function(3).output
            5

            >>> @GreatAI.create
            ... def my_function(a: int) -> int:
            ...     return a + 2
            >>> my_function(3)
            Trace[int]...

            >>> my_function('3').output
            Traceback (most recent call last):
                ...
            TypeError: type of a must be int; got str instead

        Args:
            func: The prediction function that needs to be decorated.

        Returns:
            A GreatAI instance wrapping `func`.
        """

        return GreatAI[Trace[V], V](
            func,
        )

    def __call__(self, *args: Any, **kwargs: Any) -> T:
        return self._wrapped_func(*args, **kwargs)

    @overload
    def process_batch(
        self,
        batch: Sequence[Tuple],
        *,
        concurrency: Optional[int] = None,
        unpack_arguments: Literal[True],
        do_not_persist_traces: bool = ...,
    ) -> List[Trace[V]]:
        ...

    @overload
    def process_batch(
        self,
        batch: Sequence,
        *,
        concurrency: Optional[int] = None,
        unpack_arguments: Literal[False] = ...,
        do_not_persist_traces: bool = ...,
    ) -> List[Trace[V]]:
        ...

    def process_batch(
        self,
        batch: Sequence,
        *,
        concurrency: Optional[int] = None,
        unpack_arguments: bool = False,
        do_not_persist_traces: bool = False,
    ) -> List[Trace[V]]:
        wrapped_function = self._wrapped_func

        def inner(value: Any) -> T:
            return (
                wrapped_function(*value, do_not_persist_traces=do_not_persist_traces)
                if unpack_arguments
                else wrapped_function(
                    value, do_not_persist_traces=do_not_persist_traces
                )
            )

        async def inner_async(value: Any) -> T:
            return await cast(
                Awaitable,
                (
                    wrapped_function(
                        *value, do_not_persist_traces=do_not_persist_traces
                    )
                    if unpack_arguments
                    else wrapped_function(
                        value, do_not_persist_traces=do_not_persist_traces
                    )
                ),
            )

        return list(
            tqdm(
                parallel_map(
                    inner_async
                    if get_function_metadata_store(self).is_asynchronous
                    else inner,
                    batch,
                    concurrency=concurrency,
                ),
                total=len(batch),
            )
        )

    @staticmethod
    def _get_cached_traced_function(
        func: Callable[..., Union[V, Awaitable[V]]]
    ) -> Callable[..., T]:
        @lru_cache(maxsize=get_context().prediction_cache_size)
        def func_in_tracing_context_sync(
            *args: Any,
            do_not_persist_traces: bool = False,
            **kwargs: Any,
        ) -> T:
            with TracingContext[V](
                func.__name__, do_not_persist_traces=do_not_persist_traces
            ) as t:
                result = func(*args, **kwargs)
                return cast(T, t.finalise(output=result))

        @alru_cache(maxsize=get_context().prediction_cache_size)
        async def func_in_tracing_context_async(
            *args: Any,
            do_not_persist_traces: bool = False,
            **kwargs: Any,
        ) -> T:
            with TracingContext[V](
                func.__name__, do_not_persist_traces=do_not_persist_traces
            ) as t:
                result = await cast(Callable[..., Awaitable], func)(*args, **kwargs)
                return cast(T, t.finalise(output=result))

        return cast(
            Callable[..., T],
            (
                func_in_tracing_context_async
                if get_function_metadata_store(func).is_asynchronous
                else func_in_tracing_context_sync
            ),
        )

    def _bootstrap_rest_api(
        self,
    ) -> None:
        route_config = get_context().route_config

        if route_config.prediction_endpoint_enabled:
            bootstrap_prediction_endpoint(self.app, self._wrapped_func)

        if route_config.docs_endpoints_enabled:
            bootstrap_docs_endpoints(self.app)

        if route_config.dashboard_enabled:
            bootstrap_dashboard(
                self.app,
                function_name=self.__name__,
                documentation=self.__doc__,
            )

        if route_config.trace_endpoints_enabled:
            bootstrap_trace_endpoints(self.app)

        if route_config.feedback_endpoints_enabled:
            bootstrap_feedback_endpoints(self.app)

        if route_config.meta_endpoints_enabled:
            bootstrap_meta_endpoints(
                self.app,
                self._cached_func,
                ApiMetadata(
                    name=self.__name__,
                    version=self.version,
                    documentation=self.__doc__,
                    configuration=get_context().to_flat_dict(),
                ),
            )

create(func) staticmethod #

Decorate a function by wrapping it in a GreatAI instance.

The function can be typed, synchronous or async. If it has unwrapped parameters (parameters not affected by a @parameter or @use_model decorator), those will be automatically wrapped.

The return value is replaced by a Trace (or Awaitable[Trace]), while the original return value is available under the .output property.

For configuration options, see great_ai.configure.

Examples:

>>> @GreatAI.create
... def my_function(a):
...     return a + 2
>>> my_function(3).output
5
>>> @GreatAI.create
... def my_function(a: int) -> int:
...     return a + 2
>>> my_function(3)
Trace[int]...
>>> my_function('3').output
Traceback (most recent call last):
    ...
TypeError: type of a must be int; got str instead

Parameters:

Name Type Description Default
func Union[Callable[..., Awaitable[V]], Callable[..., V]]

The prediction function that needs to be decorated.

required

Returns:

Type Description
Union[GreatAI[Awaitable[Trace[V]], V], GreatAI[Trace[V], V]]

A GreatAI instance wrapping func.

Source code in great_ai/deploy/great_ai.py
@staticmethod
def create(
    func: Union[Callable[..., Awaitable[V]], Callable[..., V]],
) -> Union["GreatAI[Awaitable[Trace[V]], V]", "GreatAI[Trace[V], V]"]:
    """Decorate a function by wrapping it in a GreatAI instance.

    The function can be typed, synchronous or async. If it has
    unwrapped parameters (parameters not affected by a @parameter
    or @use_model decorator), those will be automatically wrapped.

    The return value is replaced by a Trace (or Awaitable[Trace]),
    while the original return value is available under the `.output`
    property.

    For configuration options, see great_ai.configure.

    Examples:
        >>> @GreatAI.create
        ... def my_function(a):
        ...     return a + 2
        >>> my_function(3).output
        5

        >>> @GreatAI.create
        ... def my_function(a: int) -> int:
        ...     return a + 2
        >>> my_function(3)
        Trace[int]...

        >>> my_function('3').output
        Traceback (most recent call last):
            ...
        TypeError: type of a must be int; got str instead

    Args:
        func: The prediction function that needs to be decorated.

    Returns:
        A GreatAI instance wrapping `func`.
    """

    return GreatAI[Trace[V], V](
        func,
    )

configure(*, version='0.0.1', log_level=DEBUG, seed=42, tracing_database_factory=None, large_file_implementation=None, should_log_exception_stack=None, prediction_cache_size=512, disable_se4ml_banner=False, dashboard_table_size=20, route_config=RouteConfig()) #

Source code in great_ai/context.py
def configure(
    *,
    version: Union[int, str] = "0.0.1",
    log_level: int = DEBUG,
    seed: int = 42,
    tracing_database_factory: Optional[Type[TracingDatabaseDriver]] = None,
    large_file_implementation: Optional[Type[LargeFileBase]] = None,
    should_log_exception_stack: Optional[bool] = None,
    prediction_cache_size: int = 512,
    disable_se4ml_banner: bool = False,
    dashboard_table_size: int = 20,
    route_config: RouteConfig = RouteConfig(),
) -> None:
    global _context
    logger = get_logger("great_ai", level=log_level)

    if _context is not None:
        logger.error(
            "Configuration has been already initialised, overwriting.\n"
            + "Make sure to call `configure()` before importing your application code."
        )

    is_production = _is_in_production_mode(logger=logger)

    _set_seed(seed)

    tracing_database_factory = _initialize_tracing_database(
        tracing_database_factory, logger=logger
    )
    tracing_database = tracing_database_factory()

    if not tracing_database.is_production_ready:
        message = f"The selected tracing database ({tracing_database_factory.__name__}) is not recommended for production"
        if is_production:
            logger.error(message)
        else:
            logger.warning(message)

    _context = Context(
        version=version,
        tracing_database=tracing_database,
        large_file_implementation=_initialize_large_file(
            large_file_implementation, logger=logger
        ),
        is_production=is_production,
        logger=logger,
        should_log_exception_stack=not is_production
        if should_log_exception_stack is None
        else should_log_exception_stack,
        prediction_cache_size=prediction_cache_size,
        dashboard_table_size=dashboard_table_size,
        route_config=route_config,
    )

    logger.info(f"GreatAI (v{__version__}): configured ✅")
    for k, v in get_context().to_flat_dict().items():
        logger.info(f"{LIST_ITEM_PREFIX}{k}: {v}")

    if not is_production and not disable_se4ml_banner:
        logger.warning(
            "You still need to check whether you follow all best practices before trusting your deployment."
        )
        logger.warning(f"> Find out more at {SE4ML_WEBSITE}")

save_model(model, key, *, keep_last_n=None) #

Save (and optionally serialise) a model in order to use by use_model.

The model can be a Path or string representing a path in which case the local file/folder is read and saved using the current LargeFile implementation. In case model is an object, it is serialised using dill before uploading it.

Examples:

>>> from great_ai import use_model
>>> save_model(3, 'my_number')
'my_number:...'
>>> @use_model('my_number')
... def my_function(a, model):
...     return a + model
>>> my_function(4)
7

Parameters:

Name Type Description Default
model Union[Path, str, object]

The object or path to be uploaded.

required
key str

The model's name.

required
keep_last_n Optional[int]

If specified, remove old models and only keep the latest n. Directly passed to LargeFile.

None

Returns:

Type Description
str

The key and version of the saved model separated by a colon. Example: "key:version"

Source code in great_ai/models/save_model.py
def save_model(
    model: Union[Path, str, object], key: str, *, keep_last_n: Optional[int] = None
) -> str:
    """Save (and optionally serialise) a model in order to use by `use_model`.

    The `model` can be a Path or string representing a path in which case the
    local file/folder is read and saved using the current LargeFile implementation.
    In case `model` is an object, it is serialised using `dill` before uploading it.

    Examples:
            >>> from great_ai import use_model
            >>> save_model(3, 'my_number')
            'my_number:...'

            >>> @use_model('my_number')
            ... def my_function(a, model):
            ...     return a + model
            >>> my_function(4)
            7

    Args:
        model: The object or path to be uploaded.
        key: The model's name.
        keep_last_n: If specified, remove old models and only keep the latest n. Directly passed to LargeFile.
    Returns:
        The key and version of the saved model separated by a colon. Example: "key:version"
    """
    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}"

use_model(key, *, version='latest', model_kwarg_name='model') #

Inject a model into a function.

Load a model specified by key and version using the currently active LargeFile implementation. If it's a single object, it is deserialised using dill. If it's a directory of files, a pathlib.Path instance is given.

By default, the function's model parameter is replaced by the loaded model. This can be customised by changing model_kwarg_name. Multiple models can be loaded by decorating the same function with use_model multiple times.

Examples:

>>> from great_ai import save_model
>>> save_model(3, 'my_number')
'my_number:...'
>>> @use_model('my_number')
... def my_function(a, model):
...     return a + model
>>> my_function(4)
7
Args:
    key: The model's name as stored by the LargeFile implementation.
    version: The model's version as stored by the LargeFile implementation.
    model_kwarg_name: the parameter to use for injecting the loaded model
Returns:
    A decorator for model injection.
Source code in great_ai/models/use_model.py
def use_model(
    key: str,
    *,
    version: Union[int, Literal["latest"]] = "latest",
    model_kwarg_name: str = "model",
) -> Callable[[F], F]:
    """Inject a model into a function.

    Load a model specified by `key` and `version` using the currently active `LargeFile`
    implementation. If it's a single object, it is deserialised using `dill`. If it's a
    directory of files, a `pathlib.Path` instance is given.

    By default, the function's `model` parameter is replaced by the loaded model. This
    can be customised by changing `model_kwarg_name`. Multiple models can be loaded by
    decorating the same function with `use_model` multiple times.

    Examples:
            >>> from great_ai import save_model
            >>> save_model(3, 'my_number')
            'my_number:...'
            >>> @use_model('my_number')
            ... def my_function(a, model):
            ...     return a + model
            >>> my_function(4)
            7

        Args:
            key: The model's name as stored by the LargeFile implementation.
            version: The model's version as stored by the LargeFile implementation.
            model_kwarg_name: the parameter to use for injecting the loaded model
        Returns:
            A decorator for model injection.
    """

    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

parameter(parameter_name, *, validator=lambda _: True, disable_logging=False) #

Control the validation and logging of function parameters.

Examples:

>>> @parameter('a')
... def my_function(a: int):
...     return a + 2
>>> my_function(4)
6
>>> my_function('3')
Traceback (most recent call last):
    ...
TypeError: type of a must be int; got str instead
>>> @parameter('positive_a', validator=lambda v: v > 0)
... def my_function(positive_a: int):
...     return a + 2
>>> my_function(-1)
Traceback (most recent call last):
    ...
great_ai.errors.argument_validation_error.ArgumentValidationError: ...

Parameters:

Name Type Description Default
parameter_name str

Name of parameter to consider

required
validator Callable[[Any], bool]

Optional validator to run against the concrete argument. ArgumentValidationError is thrown when the return value is False.

lambda _: True
disable_logging bool

Do not save the value in any active TracingContext.

False

Returns:

Type Description
Callable[[F], F]

A decorator for argument validation.

Source code in great_ai/parameters/parameter.py
def parameter(
    parameter_name: str,
    *,
    validator: Callable[[Any], bool] = lambda _: True,
    disable_logging: bool = False,
) -> Callable[[F], F]:
    """Control the validation and logging of function parameters.

    Examples:
        >>> @parameter('a')
        ... def my_function(a: int):
        ...     return a + 2
        >>> my_function(4)
        6
        >>> my_function('3')
        Traceback (most recent call last):
            ...
        TypeError: type of a must be int; got str instead

        >>> @parameter('positive_a', validator=lambda v: v > 0)
        ... def my_function(positive_a: int):
        ...     return a + 2
        >>> my_function(-1)
        Traceback (most recent call last):
            ...
        great_ai.errors.argument_validation_error.ArgumentValidationError: ...

    Args:
        parameter_name: Name of parameter to consider
        validator: Optional validator to run against the concrete argument. ArgumentValidationError is thrown when the return value is False.
        disable_logging: Do not save the value in any active TracingContext.
    Returns:
        A decorator for argument validation.
    """

    def decorator(func: F) -> F:
        get_function_metadata_store(func).input_parameter_names.append(parameter_name)
        assert_function_is_not_finalised(func)

        actual_name = f"arg:{parameter_name}"

        @wraps(func)
        def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> Any:
            arguments = get_arguments(func, args, kwargs)
            argument = arguments.get(parameter_name)

            expected_type = func.__annotations__.get(parameter_name)

            if expected_type is not None:
                check_type(parameter_name, argument, expected_type)

            if not validator(argument):
                raise ArgumentValidationError(
                    f"Argument {parameter_name} in {func.__name__} did not pass validation"
                )

            context = TracingContext.get_current_tracing_context()
            if context and not disable_logging:
                context.log_value(name=f"{actual_name}:value", value=argument)
                if isinstance(argument, str):
                    context.log_value(name=f"{actual_name}:length", value=len(argument))

            return func(*args, **kwargs)

        return cast(F, wrapper)

    return decorator

log_metric(argument_name, value) #

Source code in great_ai/parameters/log_metric.py
def log_metric(argument_name: str, value: Any) -> None:
    tracing_context = TracingContext.get_current_tracing_context()
    caller = inspect.stack()[1].function
    actual_name = f"metric:{caller}:{argument_name}"
    if tracing_context:
        tracing_context.log_value(name=actual_name, value=value)

    get_context().logger.info(f"{actual_name}={value}")

Remote calls#

call_remote_great_ai #

call_remote_great_ai_async #

Ground-truth#

add_ground_truth #

query_ground_truth #

delete_ground_truth #


Last update: July 10, 2022