Improve batch processing API

This commit is contained in:
Andras Schmelczer 2022-07-08 09:36:36 +02:00
parent 1d30fd730a
commit 00714bef4e
4 changed files with 104 additions and 94 deletions

View file

@ -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(
self._wrapped_func, do_not_persist_traces=do_not_persist_traces
),
inner_async
if get_function_metadata_store(self).is_asynchronous
else inner,
batch,
concurrency=concurrency,
)

View file

@ -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,7 +32,9 @@ 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[T],
Trace( # avoid ValueError: "Trace" object has no field "__orig_class__"
trace_id=str(uuid4()),
created=self._start_time.isoformat(),
original_execution_time_ms=delta_time,
@ -49,6 +51,7 @@ class TracingContext(Generic[T]):
if get_context().is_production
else DEVELOPMENT_TAG_NAME,
],
),
)
return self._trace

View file

@ -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"]

View file

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