Support Key/Value databases

🚀 The feature

The existing cacheholder leverages a python dictionary

@functional_datapipe("in_memory_cache")
class InMemoryCacheHolderMapDataPipe(MapDataPipe[T_co]):
    def __init__(self, source_dp: MapDataPipe[T_co]) -> None:
        self.source_dp: MapDataPipe[T_co] = source_dp
        self.cache: Dict[Any, T_co] = {}

    def __getitem__(self, index) -> T_co:
        if index not in self.cache:
            self.cache[index] = self.source_dp[index] 
        return self.cache[index]  # type: ignore[index]

But could instead provide a generic interface to plug in different cache providers like redis or memcached

class Cache(ABC):
  @abstractmethod
  def __init__():
    pass
 
  @abstractmethod
  def __getitem__(self,index) -> T_co
    pass

class RedisCache(Cache):
  def __init__(self, url):
    setup_redis_client(url)
  def __getitem__(self,index) -> T_co:
    return NotImplementedError

class MemCacheCache(Cache):
  def __init__(self,url):
    setup_memcache_client(url)
  def __getitem__(self,index) -> T_co:
    return NotImplementedError

Motivation, pitch

Python dictionaries have a few limitations when used as a cache

1. Need to copy them per process
2. Updating cache in one process needs to manually synchronize with cache in other processes
3. Need to load the entire dictionary in memory to potentially look at a single element

So the goal of this work would be to reduce memory overhead and cache misses of cache in multiprocessing environments while sacrificing latency on cache hit because a python dictionary will be faster to access than a remote KV store

The 0.4 release was very much about leveraging remote object stores so this work would follow that trend

Alternatives

No response

Additional context

Our queues right now use python lists https://github.com/pytorch/data/blob/main/torchdata/dataloader2/communication/queue.py#L11 but could instead leverage queues like Kafka or RabbitMQ so imagine a similar solution

Start by reviewing the existing InMemoryCacheHolderMapDataPipe and the queue implementation in torchdata/dataloader2/communication/queue.py. Define the scope and interface for pluggable key/value cache providers such as Redis or Memcached, with completion meaning the cacheholder can use a remote provider without relying solely on a per-process dictionary.