Integrate new large_files

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"
         )