Improve batch processing API
This commit is contained in:
parent
1d30fd730a
commit
00714bef4e
4 changed files with 104 additions and 94 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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](
|
||||
trace_id=str(uuid4()),
|
||||
created=self._start_time.isoformat(),
|
||||
original_execution_time_ms=delta_time,
|
||||
logged_values=self._values,
|
||||
models=self._models,
|
||||
output=output,
|
||||
exception=None
|
||||
if exception is None
|
||||
else f"{type(exception).__name__}: {exception}",
|
||||
tags=[
|
||||
self._name,
|
||||
ONLINE_TAG_NAME,
|
||||
PRODUCTION_TAG_NAME
|
||||
if get_context().is_production
|
||||
else DEVELOPMENT_TAG_NAME,
|
||||
],
|
||||
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,
|
||||
logged_values=self._values,
|
||||
models=self._models,
|
||||
output=output,
|
||||
exception=None
|
||||
if exception is None
|
||||
else f"{type(exception).__name__}: {exception}",
|
||||
tags=[
|
||||
self._name,
|
||||
ONLINE_TAG_NAME,
|
||||
PRODUCTION_TAG_NAME
|
||||
if get_context().is_production
|
||||
else DEVELOPMENT_TAG_NAME,
|
||||
],
|
||||
),
|
||||
)
|
||||
|
||||
return self._trace
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
Loading…
Add table
Add a link
Reference in a new issue