Restore application configuration in spawned batch workers

great_ai/deploy/batch_context.py

@@ -0,0 +1,50 @@
+import copy
+import os
+from typing import Any, Dict
+
+import dill
+
+from .. import context
+from ..utilities import get_logger
+
+
+class BatchContext:
+    """Carry application configuration into independently bootstrapped workers."""
+
+    def __init__(self) -> None:
+        current = context.get_context()
+        values: Dict[str, Any] = {
+            name: getattr(current, name) for name in context.Context.model_fields
+        }
+        database = copy.copy(current.tracing_database)
+        # Built-in drivers allow configuration on their class as well as instance.
+        for name in ("path_to_db", "mongo_connection_string", "mongo_database"):
+            if hasattr(database, name):
+                setattr(database, name, getattr(database, name))
+        values["tracing_database"] = database
+        values["logger"] = (current.logger.name, current.logger.level)
+        implementation = current.large_file_implementation
+        storage = {
+            name: getattr(implementation, name)
+            for name in (
+                "cache_path",
+                "max_cache_size",
+                "initialized",
+                *implementation._configuration_attributes,
+            )
+        }
+        self._configuration = dill.dumps((values, storage))
+        self._parent_pid = os.getpid()
+        self._restored = False
+
+    def restore(self) -> None:
+        if os.getpid() == self._parent_pid or self._restored:
+            return
+        values, storage = dill.loads(self._configuration)
+        name, level = values["logger"]
+        values["logger"] = get_logger(name, level=level)
+        implementation = values["large_file_implementation"]
+        for name, value in storage.items():
+            setattr(implementation, name, value)
+        context._context = context.Context(**values)
+        self._restored = True

great_ai/deploy/great_ai.py

@@ -30,6 +30,7 @@ from ..parameters.automatically_decorate_parameters import (
 from ..tracing.tracing_context import TracingContext
 from ..utilities import parallel_map
 from ..views import ApiMetadata, Trace
+from .batch_context import BatchContext
 from .routes.bootstrap_dashboard import bootstrap_dashboard
 from .routes.bootstrap_docs_endpoints import bootstrap_docs_endpoints
 from .routes.bootstrap_feedback_endpoints import bootstrap_feedback_endpoints
@@ -210,8 +211,10 @@ class GreatAI(Generic[T, V]):
         wrapped_function = freeze_arguments(
             self._get_cached_traced_function(self._wrapped_func.__wrapped__)
         )
+        batch_context = BatchContext()
 
         def inner(value: Any) -> T:
+            batch_context.restore()
             return (
                 wrapped_function(*value, do_not_persist_traces=do_not_persist_traces)
                 if unpack_arguments
@@ -221,6 +224,7 @@ class GreatAI(Generic[T, V]):
             )
 
         async def inner_async(value: Any) -> T:
+            batch_context.restore()
             return await cast(
                 Awaitable,
                 (