[RFC] Add third party caching support in TorchServe

[agunapal]

🚀 The feature

Proposal

This proposal is for adding third party caching support in TorchServe.

Conceptually, we can add caching support in either frontend or backend

Frontend Caching Support

Backend Caching Support

From a design point of view, though it would make sense to add caching support in the frontend, the proposal is to add caching support for the backend for 2 reasons:

1. Easy of use with Python
2. Ability to extend the base class and customize the caching logic.

We start by adding support for Redis.

The implementation is based on the example proposed by @GavinPHR in #1952

Design Considerations

Two main considerations when designing this solution are

1. How do you specify the config needed to connect to RedisCache
2. Where do we plugin Caching class in TorchServe

Specify Cache Config

We could pass the config using config.properties, model-config.yaml, extra-files.
Since we are adding cache support for Python backend, it makes sense to add the config in model-config.yaml
This also gives users the flexibility to use different caches for different model.

The necessary config to access your Redis Cache is defined in the model-config.yaml

cache:
  host: localhost
  port: 6379
  db: 0

TorchServe Implementation

The basic idea is to wrap the handle method with the RedisCache class.
We use the context as the key , hash the context and use this as a key to store the inference result

Where do we wrap the handle method?

Option A

In model_loader.py, load method

User specifies config in yaml file. Based on the config, the entry_point is wrapped with RedicCache

Pro:
Minimal effort
Con:
Spaghetti code

Option B

In custom_handler.py , initialize method

Pro:
User can customize easily
Con:
User needs to define the initialize method

The below code is with Option B

To use Redis Cache with TorchServe, include the following in your handler
from ts.handler_utils.cache.redis import RedisCache

and define your handler as follows for an Image Classifier.
The key idea is to wrap the handle method with the class RedisCache

from ts.torch_handler.image_classifier import ImageClassifier


class CacheHandler(ImageClassifier):
    def __init__(self):
        super(CacheHandler, self).__init__()
        self.initialized = False

    def initialize(self, ctx):
        super().initialize(ctx)
        self.handle = RedisCache(ctx)(self.handle)
        self.initialized = True

Motivation, pitch

Implementing caching support will result in cost savings for customers in inference costs.

Alternatives

No response

Additional context

No response

[lxning]

Thanks @agunapal for the proposal. Here is my suggestions:

• define unified cache API to flexibly connect with different 3rd party caching provider.
• implement caching in handler tool library (ie. handler/util).
• how to support GPU caching?
• how to reduce latency when partial of requests can get response from caching; partial can not in a batch?
• do this caching solution support caching for a cluster of containers or only single container at local?

[agunapal]

Thanks for your suggestion @lxning

Updated Design

The Base Handler's initialize method has a decorator called cache, which takes care of handling the caching

from ts.handler.utils.cache.utils import cache

   @cache
    def initialize(self, context):

Here is the definition of the decorator

def cache(func):
    def wrap_func(self, *args, **kwargs):
        result = func(self, *args, **kwargs)
        if self.model_yaml_config is not None and "cache" in self.model_yaml_config:
            module = self.model_yaml_config["cache"]["module"]
            config = self.model_yaml_config["cache"]["config"]
            cache_class = get_cache_definition(module)
            logger.info(f"Cache used is  {cache_class}")
            self.handle = cache_class(config)(self.handle)
        return result

    return wrap_func

To take advantage of caching, users need to define the following section in their model-config.yaml

cache:
  module: <module>
  config:
     key1: val1
     key2: val2 

We can add base support for Redis and pymemcache(Client)

Redis

cache:
  module: redis
  config:
    host: localhost
    port: 6379
    db: 0

Memcached

cache:
  module: memcached
  config:
    server: '127.0.0.1:11211`

Custom Cache

If users want to use their custom solution, define a custom solution as follows for HashClient from pymemcache

import logging

from ts.handler.utils.cache.cache import Cache

logger = logging.getLogger(__name__)

try:
    import pymemcache

    _has_pymemcached = True
except ImportError:
    _has_pymemcached = False


class PyMemCached(Cache):
    def __init__(self, config=None):
        logger.info(f"Init pymemcache client with args {config}")

        if not _has_pymemcached:
            logger.error(f"Cannot import pymemcache, try pip install pymemcache.")
            return self._no_op_decorator
        self.client = pymemcache.client.hash.HashClient(**config)

        try:
            result = self.client.set('some_key', 'some_value')
        except ConnectionRefusedError:
            logger.error(
                f"Cannot connect to a memcached server, ensure a server is running on {config['server']}."
            )

This file needs to be passed an --extra-file with torch-model-archiver and the file needs to be specified in the model-config.yaml as shown below

Example:

• Custom

cache:
  module: custom.py
  config:
    servers: ['127.0.0.1:11211']

torch-model-archiver --model-name alexnet --version 1.0 --model-file model.py --serialized-file alexnet-owt-7be5be79.pth --handler image_classifier --config-file model-config.yaml --extra-files=/home/ubuntu/serve/examples/caching/custom.py

Miscellaneous

• The solution works on both CPU & GPU
• The solution would work best for single requests. In case of batched requests, the latency might be more

[andrewcatalano]

Hi @agunapal , this looks great, really comprehensive explanation as well. Just as an fyi for pymemcached, at least for our use case, we’re going to connect to memcache using the glb and will just have one client - in this case the benefits of using HashClient go away. Going to see if there are any improvements using the PooledClient. More importantly, we’ve observed some queue latency when making sets with an unseen dataset at higher TPS. Pymemcached doesn’t support async sets - have you considered putting an async wrapper on this? I’ll be working on an implementation of this today and will let you know if we see any benefit.