Skip to content

feat: add caching to GapicCallable #666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 35 commits into from
Oct 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
35 commits
Select commit Hold shift + click to select a range
5518645
feat: optimize _GapicCallable
daniel-sanche Sep 6, 2023
0cc03bd
cleaned up metadata lines
daniel-sanche Sep 7, 2023
c97a636
chore: avoid type checks in error wrapper
daniel-sanche Sep 7, 2023
b453df4
Revert "chore: avoid type checks in error wrapper"
daniel-sanche Sep 7, 2023
2f7acff
add default wrapped function
daniel-sanche Sep 8, 2023
31f0b4e
fixed decorator order
daniel-sanche Sep 8, 2023
b92328c
fixed spacing
daniel-sanche Sep 8, 2023
0831dbf
fixed comment typo
daniel-sanche Sep 8, 2023
a1563d2
fixed spacing
daniel-sanche Sep 8, 2023
fb1a372
Merge branch 'main' into optimize_gapic_callable
daniel-sanche Sep 11, 2023
52ed5be
Merge branch 'main' into optimize_gapic_callable
daniel-sanche Feb 10, 2024
85e2102
fixed spacing
daniel-sanche Feb 10, 2024
c76f51c
removed unneeded helpers
daniel-sanche Feb 10, 2024
f4a9021
use caching
daniel-sanche Feb 10, 2024
cacc73c
improved metadata parsing
daniel-sanche Feb 10, 2024
a30101d
improved docstring
daniel-sanche Feb 10, 2024
db9a9c4
fixed logic
daniel-sanche Feb 10, 2024
a555629
added benchmark test
daniel-sanche Feb 13, 2024
cfe5c7d
Merge branch 'main' into optimize_gapic_callable
daniel-sanche Feb 13, 2024
fbbaaca
update threshold
daniel-sanche Feb 15, 2024
576bb0f
run benchmark in loop for testing
daniel-sanche Feb 16, 2024
7c32a5d
use verbose logs
daniel-sanche Feb 16, 2024
26bec79
Revert testing
daniel-sanche Feb 16, 2024
c25e0eb
used smaller value
daniel-sanche Feb 16, 2024
49201ca
changed threshold
daniel-sanche Feb 16, 2024
3d2c964
Merge branch 'main' into optimize_gapic_callable
daniel-sanche Feb 16, 2024
c9c1ff3
Merge branch 'main' into optimize_gapic_callable
daniel-sanche May 3, 2024
deca58c
removed link in comment
daniel-sanche May 3, 2024
9bf72f5
use list type for metadata
daniel-sanche Jun 4, 2024
3f9066b
add types to docstring
daniel-sanche Jun 4, 2024
45858b8
Merge branch 'main' into optimize_gapic_callable
daniel-sanche Jun 4, 2024
fd37981
fixed lint
daniel-sanche Jun 4, 2024
f08fb98
convert to list at init time
daniel-sanche Jun 4, 2024
4c6e94c
Merge branch 'main' into optimize_gapic_callable
parthea Oct 21, 2024
bcd019b
Merge branch 'main' into optimize_gapic_callable
parthea Oct 21, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 36 additions & 44 deletions google/api_core/gapic_v1/method.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,24 +42,6 @@ class _MethodDefault(enum.Enum):
so the default should be used."""


def _is_not_none_or_false(value):
return value is not None and value is not False


def _apply_decorators(func, decorators):
"""Apply a list of decorators to a given function.

``decorators`` may contain items that are ``None`` or ``False`` which will
be ignored.
"""
filtered_decorators = filter(_is_not_none_or_false, reversed(decorators))

for decorator in filtered_decorators:
func = decorator(func)

return func


class _GapicCallable(object):
"""Callable that applies retry, timeout, and metadata logic.

Expand Down Expand Up @@ -91,44 +73,53 @@ def __init__(
):
self._target = target
self._retry = retry
if isinstance(timeout, (int, float)):
timeout = TimeToDeadlineTimeout(timeout=timeout)
self._timeout = timeout
self._compression = compression
self._metadata = metadata
self._metadata = list(metadata) if metadata is not None else None

def __call__(
self, *args, timeout=DEFAULT, retry=DEFAULT, compression=DEFAULT, **kwargs
):
"""Invoke the low-level RPC with retry, timeout, compression, and metadata."""

if retry is DEFAULT:
retry = self._retry

if timeout is DEFAULT:
timeout = self._timeout

if compression is DEFAULT:
compression = self._compression

if isinstance(timeout, (int, float)):
timeout = TimeToDeadlineTimeout(timeout=timeout)

# Apply all applicable decorators.
wrapped_func = _apply_decorators(self._target, [retry, timeout])
if compression is not None:
kwargs["compression"] = compression

# Add the user agent metadata to the call.
if self._metadata is not None:
metadata = kwargs.get("metadata", [])
# Due to the nature of invocation, None should be treated the same
# as not specified.
if metadata is None:
metadata = []
metadata = list(metadata)
metadata.extend(self._metadata)
kwargs["metadata"] = metadata
if self._compression is not None:
kwargs["compression"] = compression
try:
# attempt to concatenate default metadata with user-provided metadata
kwargs["metadata"] = [*kwargs["metadata"], *self._metadata]
except (KeyError, TypeError):
# if metadata is not provided, use just the default metadata
kwargs["metadata"] = self._metadata

call = self._build_wrapped_call(timeout, retry)
return call(*args, **kwargs)

@functools.lru_cache(maxsize=4)
def _build_wrapped_call(self, timeout, retry):
"""
Build a wrapped callable that applies retry, timeout, and metadata logic.
"""
wrapped_func = self._target
if timeout is DEFAULT:
timeout = self._timeout
elif isinstance(timeout, (int, float)):
timeout = TimeToDeadlineTimeout(timeout=timeout)
if timeout is not None:
wrapped_func = timeout(wrapped_func)

if retry is DEFAULT:
retry = self._retry
if retry is not None:
wrapped_func = retry(wrapped_func)

return wrapped_func(*args, **kwargs)
return wrapped_func


def wrap_method(
Expand Down Expand Up @@ -202,8 +193,9 @@ def get_topic(name, timeout=None):

Args:
func (Callable[Any]): The function to wrap. It should accept an
optional ``timeout`` argument. If ``metadata`` is not ``None``, it
should accept a ``metadata`` argument.
optional ``timeout`` (google.api_core.timeout.Timeout) argument.
If ``metadata`` is not ``None``, it should accept a ``metadata``
(Sequence[Tuple[str, str]]) argument.
default_retry (Optional[google.api_core.Retry]): The default retry
strategy. If ``None``, the method will not retry by default.
default_timeout (Optional[google.api_core.Timeout]): The default
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/gapic/test_method.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,3 +222,24 @@ def test_wrap_method_with_call_not_supported():
with pytest.raises(ValueError) as exc_info:
google.api_core.gapic_v1.method.wrap_method(method, with_call=True)
assert "with_call=True is only supported for unary calls" in str(exc_info.value)


def test_benchmark_gapic_call():
"""
Ensure the __call__ method performance does not regress from expectations

__call__ builds a new wrapped function using passed-in timeout and retry, but
subsequent calls are cached

Note: The threshold has been tuned for the CI workers. Test may flake on
slower hardware
"""
from google.api_core.gapic_v1.method import _GapicCallable
from google.api_core.retry import Retry
from timeit import timeit

gapic_callable = _GapicCallable(
lambda *a, **k: 1, retry=Retry(), timeout=1010, compression=False
)
avg_time = timeit(lambda: gapic_callable(), number=10_000)
assert avg_time < 0.4