great-ai/great_ai/context.py

import os
import random
from logging import DEBUG, Logger
from pathlib import Path
from typing import Any, Dict, Optional, Type, cast

from pydantic import BaseModel

from .constants import (
    DEFAULT_LARGE_FILE_CONFIG_PATHS,
    DEFAULT_TRACING_DATABASE_CONFIG_PATHS,
    ENV_VAR_KEY,
    LIST_ITEM_PREFIX,
    PRODUCTION_KEY,
    SE4ML_WEBSITE,
)
from .large_file import LargeFileBase, LargeFileLocal
from .persistence import ParallelTinyDbDriver, TracingDatabaseDriver
from .utilities import get_logger


class Context(BaseModel):
    tracing_database: TracingDatabaseDriver
    large_file_implementation: Type[LargeFileBase]
    is_production: bool
    logger: Logger
    should_log_exception_stack: bool
    prediction_cache_size: int
    dashboard_table_size: int

    class Config:
        arbitrary_types_allowed = True

    def to_flat_dict(self) -> Dict[str, Any]:
        return {
            "tracing_database": type(self.tracing_database).__name__,
            "large_file_implementation": self.large_file_implementation.__name__,
            "is_production": self.is_production,
            "should_log_exception_stack": self.should_log_exception_stack,
            "prediction_cache_size": self.prediction_cache_size,
            "dashboard_table_size": self.dashboard_table_size,
        }


_context: Optional[Context] = None


def get_context() -> Context:
    if _context is None:
        configure()

    return cast(Context, _context)


def configure(
    *,
    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,
) -> None:
    global _context
    logger = get_logger("great_ai", level=log_level)

    if _context is not None:
        logger.warn(
            "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(
        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,
    )

    logger.info("Settings: 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}")


def _is_in_production_mode(logger: Optional[Logger]) -> bool:
    environment = os.environ.get(ENV_VAR_KEY)

    if environment is None:
        if logger:
            logger.warning(
                f"Environment variable {ENV_VAR_KEY} is not set, defaulting to development mode ‼️"
            )
        is_production = False
    else:
        is_production = environment.lower() == PRODUCTION_KEY
        if logger:
            if not is_production:
                logger.info(
                    f"Value of {ENV_VAR_KEY} is `{environment}` which is not equal to `{PRODUCTION_KEY}`"
                    + " defaulting to development mode ‼️"
                )
            else:
                logger.info("Running in production mode ✅")

    return is_production


def _initialize_tracing_database(
    selected: Optional[Type[TracingDatabaseDriver]], logger: Logger
) -> Type[TracingDatabaseDriver]:
    for tracing_driver, paths in DEFAULT_TRACING_DATABASE_CONFIG_PATHS.items():
        if selected is None or selected == tracing_driver:
            if tracing_driver.initialized:
                logger.warning(
                    f"{tracing_driver.__name__} has been already configured: skipping initialisation"
                )
                return tracing_driver
            for p in paths:
                if Path(p).exists():
                    logger.info(
                        f"Found credentials file ({Path(p).absolute()}), initialising {tracing_driver.__name__}"
                    )
                    tracing_driver.configure_credentials_from_file(p)
                    return tracing_driver
    logger.warning(
        "Cannot find credentials files, defaulting to using ParallelTinyDbDriver"
    )
    return ParallelTinyDbDriver


def _initialize_large_file(
    selected: Optional[Type[LargeFileBase]], logger: Logger
) -> Type[LargeFileBase]:
    for large_file, paths in DEFAULT_LARGE_FILE_CONFIG_PATHS.items():
        if selected is None or selected == large_file:
            if large_file.initialized:
                logger.warning(
                    f"{large_file.__name__} has been already configured: skipping initialisation"
                )
                return large_file
            for p in paths:
                if Path(p).exists():
                    logger.info(
                        f"Found credentials file ({Path(p).absolute()}), initialising {large_file.__name__}"
                    )
                    large_file.configure_credentials_from_file(p)
                    return large_file
    logger.warning("Cannot find credentials files, defaulting to using LargeFileLocal")
    return LargeFileLocal


def _set_seed(seed: int) -> None:
    random.seed(seed)

    try:
        import numpy

        numpy.random.seed(seed + 1)
    except ImportError:
        pass