Improve batch processing API

great_ai/deploy/great_ai.py

@@ -1,4 +1,4 @@
-from functools import lru_cache, partial, wraps
+from functools import lru_cache, wraps
 from textwrap import dedent
 from typing import (
     Any,
@@ -6,8 +6,10 @@ from typing import (
     Callable,
     Generic,
     List,
+    Literal,
     Optional,
     Sequence,
+    Tuple,
     TypeVar,
     Union,
     cast,
@@ -96,17 +98,66 @@ class GreatAI(Generic[T, V]):
     def __call__(self, *args: Any, **kwargs: Any) -> T:
         return self._wrapped_func(*args, **kwargs)
 
+    @overload
+    def process_batch(
+        self,
+        batch: Sequence[Tuple],
+        *,
+        concurrency: Optional[int] = None,
+        unpack_arguments: Literal[True],
+        do_not_persist_traces: bool = ...,
+    ) -> List[Trace[V]]:
+        ...
+
+    @overload
     def process_batch(
         self,
         batch: Sequence,
+        *,
         concurrency: Optional[int] = None,
-        do_not_persist_traces: Optional[bool] = False,
+        unpack_arguments: Literal[False] = ...,
+        do_not_persist_traces: bool = ...,
     ) -> List[Trace[V]]:
+        ...
+
+    def process_batch(
+        self,
+        batch: Sequence,
+        *,
+        concurrency: Optional[int] = None,
+        unpack_arguments: bool = False,
+        do_not_persist_traces: bool = False,
+    ) -> List[Trace[V]]:
+        wrapped_function = self._wrapped_func
+
+        def inner(value: Any) -> T:
+            return (
+                wrapped_function(*value, do_not_persist_traces=do_not_persist_traces)
+                if unpack_arguments
+                else wrapped_function(
+                    value, do_not_persist_traces=do_not_persist_traces
+                )
+            )
+
+        async def inner_async(value: Any) -> T:
+            return await cast(
+                Awaitable,
+                (
+                    wrapped_function(
+                        *value, do_not_persist_traces=do_not_persist_traces
+                    )
+                    if unpack_arguments
+                    else wrapped_function(
+                        value, do_not_persist_traces=do_not_persist_traces
+                    )
+                ),
+            )
+
         return list(
             parallel_map(
-                partial(
+                inner_async
-                    self._wrapped_func, do_not_persist_traces=do_not_persist_traces
+                if get_function_metadata_store(self).is_asynchronous
-                ),
+                else inner,
                 batch,
                 concurrency=concurrency,
             )

great_ai/tracing/tracing_context.py

@@ -1,7 +1,7 @@
 from contextvars import ContextVar
 from datetime import datetime
 from types import TracebackType
-from typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar
+from typing import Any, Dict, Generic, List, Literal, Optional, Type, TypeVar, cast
 from uuid import uuid4
 
 from ..constants import DEVELOPMENT_TAG_NAME, ONLINE_TAG_NAME, PRODUCTION_TAG_NAME
@@ -32,23 +32,26 @@ class TracingContext(Generic[T]):
         assert self._trace is None, "has been already finalised"
 
         delta_time = (datetime.utcnow() - self._start_time).microseconds / 1000
-        self._trace = Trace[T](
+        self._trace = cast(
-            trace_id=str(uuid4()),
+            Trace[T],
-            created=self._start_time.isoformat(),
+            Trace(  # avoid ValueError: "Trace" object has no field "__orig_class__"
-            original_execution_time_ms=delta_time,
+                trace_id=str(uuid4()),
-            logged_values=self._values,
+                created=self._start_time.isoformat(),
-            models=self._models,
+                original_execution_time_ms=delta_time,
-            output=output,
+                logged_values=self._values,
-            exception=None
+                models=self._models,
-            if exception is None
+                output=output,
-            else f"{type(exception).__name__}: {exception}",
+                exception=None
-            tags=[
+                if exception is None
-                self._name,
+                else f"{type(exception).__name__}: {exception}",
-                ONLINE_TAG_NAME,
+                tags=[
-                PRODUCTION_TAG_NAME
+                    self._name,
-                if get_context().is_production
+                    ONLINE_TAG_NAME,
-                else DEVELOPMENT_TAG_NAME,
+                    PRODUCTION_TAG_NAME
-            ],
+                    if get_context().is_production
+                    else DEVELOPMENT_TAG_NAME,
+                ],
+            ),
         )
 
         return self._trace

tests/test_great_ai.py

@@ -12,6 +12,19 @@ def test_process_batch() -> None:
     assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36]
 
 
+def test_process_batch_unpacked() -> None:
+    @GreatAI.create
+    def f(a: int, b: str) -> str:
+        return b * a
+
+    assert [
+        v.output
+        for v in f.process_batch(
+            [(2, "aa"), (1, "fa"), (3, "b")], unpack_arguments=True
+        )
+    ] == ["aaaa", "fa", "bbb"]
+
+
 @pytest.mark.asyncio
 async def test_process_batch_async() -> None:
     @GreatAI.create
@@ -20,3 +33,17 @@ async def test_process_batch_async() -> None:
         return x + 2
 
     assert [v.output for v in f.process_batch([3, 9, 34])] == [5, 11, 36]
+
+
+@pytest.mark.asyncio
+async def test_process_batch_async_unpacked() -> None:
+    @GreatAI.create
+    async def f(a: int, b: str) -> str:
+        return b * a
+
+    assert [
+        v.output
+        for v in f.process_batch(
+            [(2, "aa"), (1, "fa"), (3, "b")], unpack_arguments=True
+        )
+    ] == ["aaaa", "fa", "bbb"]

tests/utilities/test_match_names.py

@@ -1,71 +0,0 @@
-import unittest
-
-from great_ai.utilities import match_names
-
-
-class TestMatchNames(unittest.TestCase):
-    def test_grid(self) -> None:
-        names = [
-            ["Al-Nasiry, S.", "Al-Nasiry S"],
-            ["De Leo, A. (Andreina)"],
-            [
-                "DeBenedictis, J.N. (Julia)",
-                "DeBenedictis, N. (Julia)",
-                "DeBenedictis, J.",
-                "Julia DeBenedictis",
-            ],
-            ["Diesen, J.A.Y. van"],
-            ["Nio, C Yung"],
-            ["Uyen Chau Nguyen"],
-            ["Hagger M.S.", "Martin Hagger"],
-            ["Izabela-Cristina STANCU"],
-            ["Verborgh R.", "Ruben Verborgh"],
-            ["Droshout, D.T.G.G.M.L. (Dimitri)"],
-            ["El Demellawy"],
-            ["Sebastiaan van de Velde", "van de Velde, Sebastiaan"],
-            ["Sebastiaan Brand"],
-            ["Bertil FM Blok", "B.F.M. Blok"],
-            ["Bob Zadok Blok"],
-            ["Shannon Spruit"],
-            ["Shannon Kroes"],
-            [
-                "MSc Jérémie Decouchant",
-                "PhD, Jérémie Decouchant",
-                "PhD Jérémie Decouchant",
-                "Jérémie Decouchant",
-                "MD, PhD Jérémie Decouchant",
-            ],
-            ["Jeremie Gobeil"],
-            ["Wouters, B.B.R.E.F.", "Wouters B.B. R.  E.F."],
-            ["Elias Caldeira Dantas, A.M. (Aline)"],
-            ["Ed Harris"],
-            ["Ed Deprettere ", "    Ed   Deprettere "],
-            ["András Schmelczer"],
-            ["Enden, G. van den (Gitta)"],
-            ["Ruben van Dijk"],
-            ["Richard van Dijk"],
-            ["Xinping Guan", "X. Guan"],
-            ["Xiaohong Guan", "X. Guan"],
-            ["Pruimboom, Tim", "Pruimboom T."],
-            ["Sanchez-Faddeev H.", "Sanchez-Faddeev, Hernando"],
-            ["Duijnhoven - Jansen, E.M. van", "Emma M. van Jansen"],
-        ]
-
-        all_names = [n for t in names for n in t]
-
-        for n1 in all_names:
-            for n2 in all_names:
-                is_match = match_names(n1, n2)
-                groups = [t for t in names if n1 in t]
-                with self.subTest(n1=n1, n2=n2):
-                    if any(n2 in g for g in groups):
-                        self.assertTrue(is_match, "false negative")
-                    elif all(n2 not in g for g in groups):
-                        self.assertFalse(is_match, "false match")
-
-    def test_empty(self) -> None:
-        self.assertFalse(match_names("", ""))
-        self.assertFalse(match_names(None, ""))
-        self.assertFalse(match_names(None, None))
-        self.assertFalse(match_names("", None))
-        self.assertFalse(match_names("Oliver", None))