From 00714bef4ecfa16217bbffd903a0c4845cb6fcc3 Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Fri, 8 Jul 2022 09:36:36 +0200 Subject: [PATCH] Improve batch processing API --- great_ai/deploy/great_ai.py | 61 +++++++++++++++++++++++-- great_ai/tracing/tracing_context.py | 39 ++++++++-------- tests/test_great_ai.py | 27 +++++++++++ tests/utilities/test_match_names.py | 71 ----------------------------- 4 files changed, 104 insertions(+), 94 deletions(-) delete mode 100644 tests/utilities/test_match_names.py diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py index a3059f7..a3c3b35 100644 --- a/great_ai/deploy/great_ai.py +++ b/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( - 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, ) diff --git a/great_ai/tracing/tracing_context.py b/great_ai/tracing/tracing_context.py index 167132b..35d298e 100644 --- a/great_ai/tracing/tracing_context.py +++ b/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]( - 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 diff --git a/tests/test_great_ai.py b/tests/test_great_ai.py index 17e5e65..60ee511 100644 --- a/tests/test_great_ai.py +++ b/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"] diff --git a/tests/utilities/test_match_names.py b/tests/utilities/test_match_names.py deleted file mode 100644 index bd66ac0..0000000 --- a/tests/utilities/test_match_names.py +++ /dev/null @@ -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))