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,
                 (

tests/test_batch_configuration.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+
+import pytest
+from great_ai import GreatAI, ParallelTinyDbDriver, RouteConfig, configure
+from great_ai.context import get_context
+
+
+def test_spawn_batch_uses_parent_configuration(
+    isolated_context: None, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    database = tmp_path / "custom-traces.json"
+    monkeypatch.setattr(ParallelTinyDbDriver, "path_to_db", database)
+    configure(
+        prediction_cache_size=7,
+        disable_se4ml_banner=True,
+        route_config=RouteConfig(dashboard_enabled=False),
+    )
+
+    @GreatAI.create
+    def predict(value: int) -> int:
+        return get_context().prediction_cache_size + value
+
+    monkeypatch.setattr("multiprocessing.get_all_start_methods", lambda: ["spawn"])
+    assert [t.output for t in predict.process_batch([1, 2], concurrency=1)] == [8, 9]
+    assert database.exists()
+    assert len(get_context().tracing_database.query()[0]) == 2