Integrate new large_files

great_ai/src/great_ai/great_ai/constants.py

@@ -1,4 +1,14 @@
+from pathlib import Path
+
+from great_ai.large_file import LargeFileLocal, LargeFileMongo, LargeFileS3
+
 ENV_VAR_KEY = "ENVIRONMENT"
 PRODUCTION_KEY = "production"
 DEFAULT_TRACING_DB_FILENAME = "tracing_database.json"
 METRICS_PATH = "/metrics"
+
+DEFAULT_LARGE_FILE_CONFIG_PATHS = {
+    LargeFileLocal: None,
+    LargeFileMongo: Path("mongodb.ini"),
+    LargeFileS3: Path("s3.ini"),
+}

great_ai/src/great_ai/great_ai/context/configure.py

@@ -2,24 +2,29 @@ import os
 import random
 from logging import DEBUG, Logger
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Type
 
 import great_ai.great_ai.context.context as context
-from great_ai.open_s3 import LargeFile
+from great_ai.large_file import LargeFile, LargeFileLocal
 from great_ai.utilities.logger import get_logger
 
-from ..constants import DEFAULT_TRACING_DB_FILENAME, ENV_VAR_KEY, PRODUCTION_KEY
-from ..persistence import ParallelTinyDbDriver, PersistenceDriver
+from ..constants import (
+    DEFAULT_LARGE_FILE_CONFIG_PATHS,
+    DEFAULT_TRACING_DB_FILENAME,
+    ENV_VAR_KEY,
+    PRODUCTION_KEY,
+)
+from ..tracing.parallel_tinydb_driver import ParallelTinyDbDriver, TracingDatabase
 
 
 def configure(
     version: str = "0.0.1",
     log_level: int = DEBUG,
-    s3_config: Path = Path("s3.ini"),
     seed: int = 42,
-    persistence_driver: PersistenceDriver = ParallelTinyDbDriver(
+    tracing_database: TracingDatabase = ParallelTinyDbDriver(
         Path(DEFAULT_TRACING_DB_FILENAME)
     ),
+    large_file_implementation: Type[LargeFile] = LargeFileLocal,
     should_log_exception_stack: Optional[bool] = None,
     prediction_cache_size: int = 512,
 ) -> None:
@@ -28,21 +33,22 @@ def configure(
     if context._context is not None:
         logger.warn(
             "Configuration has been already initialised, overwriting.\n"
-            + 'Make sure to call "configure()" before importing your application code.'
+            + "Make sure to call `configure()` before importing your application code."
         )
 
     is_production = _is_in_production_mode(logger=logger)
-    _initialize_large_file(s3_config, logger=logger)
+    _initialize_large_file(large_file_implementation, logger=logger)
     _set_seed(seed)
 
-    if not persistence_driver.is_threadsafe:
+    if not tracing_database.is_threadsafe:
         logger.warning(
-            f"The selected persistence driver ({type(persistence_driver).__name__}) is not threadsafe"
+            f"The selected persistence driver ({type(tracing_database).__name__}) is not threadsafe"
         )
 
     context._context = context.Context(
         version=version,
-        persistence=persistence_driver,
+        tracing_database=tracing_database,
+        large_file_implementation=large_file_implementation,
         is_production=is_production,
         logger=logger,
         should_log_exception_stack=not is_production
@@ -77,12 +83,23 @@ def _is_in_production_mode(logger: Optional[Logger]) -> bool:
     return is_production
 
 
-def _initialize_large_file(s3_config: Path, logger: Logger) -> None:
-    if s3_config.exists():
-        LargeFile.configure_credentials_from_file(s3_config)
+def _initialize_large_file(large_file: Type[LargeFile], logger: Logger) -> None:
+    path = DEFAULT_LARGE_FILE_CONFIG_PATHS[large_file]
+    if path is None:
+        return
+
+    if large_file.initialized:
+        logger.warning(
+            f"{large_file.__name__} has been already configured: skipping initialisation"
+        )
+        return
+
+    if path.exists():
+        large_file.configure_credentials_from_file(path)
+        logger.info(f"{large_file.__name__} initialised with config ({path.resolve()})")
     else:
         logger.warning(
-            f"Provided S3 config ({s3_config.resolve()}) not found, skipping LargeFile initialisation"
+            f"Default {large_file.__name__} config ({path.resolve()}) not found, skipping {large_file.__name__} initialisation"
         )
 
 

great_ai/src/great_ai/great_ai/context/context.py

@@ -1,14 +1,17 @@
 from logging import Logger
-from typing import Optional
+from typing import Any, Dict, Optional, Type
 
 from pydantic import BaseModel
 
-from ..persistence import PersistenceDriver
+from great_ai.large_file.large_file.large_file import LargeFile
+
+from ..tracing.tracing_database import TracingDatabase
 
 
 class Context(BaseModel):
     version: str
-    persistence: PersistenceDriver
+    tracing_database: TracingDatabase
+    large_file_implementation: Type[LargeFile]
     is_production: bool
     logger: Logger
     should_log_exception_stack: bool
@@ -17,5 +20,15 @@ class Context(BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
+    def to_flat_dict(self) -> Dict[str, Any]:
+        return {
+            "version": self.version,
+            "tracing_database": type(self.tracing_database).__name__,
+            "large_file_implementation": self.large_file_implementation.__name__,
+            "is_production": self.is_production,
+            "logger": type(self.logger).__name__,
+            "should_log_exception_stack": self.should_log_exception_stack,
+        }
+
 
 _context: Optional[Context] = None