[Roadmap] A general storage abstraction layer

[0oshowero0]

Background

In our previous implementation, we provided a naive TransferQueueStorageSimpleUnit class as the storage layer, which is tightly coupled with other components such as TransferQueueController and TransferQueueClient. This introduces extra code logic that may not be universal, making it difficult to support high-performance storage backends such as MoonCakeStore and Ray Direct Transport.

For instance,

• In TransferQueueController: The controller manages local_indexes for SimpleUnit, which is not universal information for all storage backends.
• In TransferQueueClient: The client directly manages the distributed SimpleUnit, making it difficult to adapt to other backends without major refactoring.

Solution

In PR66, we propose a general abstraction class TransferQueueStorageManager for storage backends. It defines the following interfaces:

    @abstractmethod
    async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None:
        raise NotImplementedError("Subclasses must implement put_data")

    @abstractmethod
    async def get_data(self, metadata: BatchMeta) -> TensorDict:
        raise NotImplementedError("Subclasses must implement get_data")

    @abstractmethod
    async def clear_data(self, metadata: BatchMeta) -> None:
        raise NotImplementedError("Subclasses must implement clear_data")

For customized storage backends, you only need to inherit from TransferQueueStorageManager, and implement these functions according to the API of your storage backends.

For general functions used in TransferQueue, such as _do_handshake_with_controller and notify_data_update, we also include them in this Manager class.

---

We will provide a KVStorageManager for general KV based storage as a reference implementation, where we can integrate different backends through simple configuration.

---

In this PR, we have rewritten the TransferQueueStorageSimpleUnit backend according to the above logic. We implement AsyncSimpleStorageManager to hide the complex logic of manipulating distributed storage units from the user interface. Now in TransferQueueClient, the core function for async_get_data looks like this:

# In class AsyncTransferQueueClient
    async def async_get_data(self, metadata: BatchMeta) -> TensorDict:

        if not metadata or metadata.size == 0:
            return TensorDict({}, batch_size=0)

        results = await self.storage_manager.get_data(metadata)

        return results

The self.storage_manager is initialized by a factory:

# In class AsyncTransferQueueClient
    def initialize_storage_manager(
        self,
        manager_type: str,
        config: dict[str, Any],
    ):
        self.storage_manager = TransferQueueStorageManagerFactory.create(manager_type, config)

Progress

We have completed the initial refactoring and provided an E2E script in PR66, and it's now available in dev branch. You can try it by running the following commands:

python recipe/simple_use_case/async_demo.py

TODO

• Code optimization & docstrings @LLLLxmmm @jianjunzhong
• Provide a general dynamic socket function for both Client & Storage @ji-huazhong.
• Fix all UTs @LLLLxmmm @jianjunzhong @0oshowero0
• Provide sync_demo @LLLLxmmm @jianjunzhong (In following PR)
• Provide a KVStorageManager implementation, where we can select different backend storages like Mooncake, Redis, etc. (PR#96)
• Provide a general data system initialization utility function (In following PR)
• Better folder structure

[0oshowero0]

The KVStorageManager may look like this:

# We hope no modification is needed for this class when adding new KV-based storage backends
class KVStorageManger(TransferQueueStorageManager):
    def __init__(self, config):
        super().__init__(config)

        self.client = KVClientFactory(config)

    def get_data(self, BatchMeta):
        ...
        data = self.client.get(key, optional_infos)   # the get/put function for different backends should be similar
        ...

        batch_data = concat(data)
        return batch_data
    
    def put_data():
        pass
    
    def clean_data():
        pass


# A thin layer inheriting from KVStorageManager to handle backend-specific configurations
class MoonCakeStorageManger(KVStroageManger):
    def __init__(self, config):
        # Validate configuration for MoonCake backend
        self._validate_config(config)
        super().__init__(config)

    def _validate_config(self, config):
        pass

# A factory class for initializing different KV storage clients
class KVClientFactory:
    _registry: dict[str, Any] = {}

    @classmethod
    def register(cls, kv_client_type: str, kv_client_cls):
        pass
        
    @classmethod
    def create(cls, kv_client_type: str, config: dict[str, Any]):
        pass

[0oshowero0]

To support various storage backends, we plan to organize the storage folder as follows:

  transfer_queue/
  ├── storage/
  │   ├── __init__.py              
  │   │── simple_backend.py             # SimpleStorageUnit、StorageUnitData、StorageMetaGroup
  │   ├── managers/
  │   │   ├── __init__.py
  |   |   |──base.py                    # TransferQueueStorageManager, KVStorageManager
  │   │   ├──simple_backend_manager.py  # AsyncSimpleStorageManager
  │   │   ├──yuanrong_manager.py        # YuanRongStorageManager
  │   │   ├──mooncake_manager.py        # MoonCakeStorageManager
  │   │   └──factory.py                 # TransferQueueStorageManagerFactory
  │   └── clients/
  │   │   ├── __init__.py
  │   │   └── yuanrong_client.py
  │   │   └── mooncake_client.py
  │   │   └── factory.py                # TransferQueueStorageClientFactory