Refactor storage service to increase efficiency and stability

Nobody has claimed this yet.

# Motivation

Many problems exist with current implementation of Mars storage.

1. No flexible way to control data location

When loading data from other endpoints, we may prefer multiple locations for fallback. Current implementation does not support this and may introduce unnecessary spill operations.

2. No support for remote reader / writer

Remote readers and writers provide flexible way to handle data transfer, enabling shuffle and client-side data manipulation without high memory cost. Current implementation only handles readers and writers locally.

3. Mix of lower-level code and higher-level code

Data transfer and spill should be implemented upon a common IO layer to make the whole storage more maintainable. Current implementation mixes all things up.

4. Race condition exists when spilling data on shuffle

In current implementation, when starting a reader and data spill is launched, it is possible that the data is spilled and we get a KeyError afterward.

5. Unnecessary IPC calls

In current implementation, we need to do quota request, put data info, update quota and deal with spill, all introducing more than one IPC call. The number of calls can be reduced to no more than 2.

# Design

The new design of Mars storage can be divided into two parts: the kernel storage and the user storage. The kernel storage is a thin wrap of storage backends plus necessary access controls. The user storage is constructed over the kernel storage with spill and transfer support.

## Kernel Storage

The principle of kernel storage is to make thing simple. That is, the API does not handle complicated retries and redirections. When encountering storage full or lock errors, it raises straightforwardly (instead of performing wait or retry). `KernelStorageAPI` will look like

```python
class KernelStorageAPI:
@classmethod
async def create(cls, band_name: str, worker_address: str) -> "KernelStorageAPI":
"""
Create a band-specific API
"""

async def open_reader(
self,
session_id: str,
data_key: str,
level: StorageLevel = None,
) -> KernelStorageFileObject:
"""
Create a reader on a specific file
"""

async def open_writer(
self,
session_id: str,
data_key: str,
size: int,
level: StorageLevel = None,
) -> KernelStorageFileObject:
"""
Create a writer on a specific file
"""

async def delete(self, session_id: str, data_key: str, error: str = "raise"):
"""
Delete a file with specified keys
"""

async def get_capacity(self) -> Dict[StorageLevel, StorageCapacity]:
"""
Get capacities of levels of the band
"""

async def list(
self,
level: StorageLevel,
lock_free_only: bool = False,
) -> List[InternalDataInfo]:
"""
Get information of all data in the band
"""

async def put(
self,
session_id: str,
data_key: str,
obj: Any,
level: StorageLevel = None,
) -> InternalDataInfo:
"""
Put an object into the band storage
"""

async def get(
self,
session_id: str,
data_key: str,
conditions: List = None,
level: StorageLevel = None,
error: str = "raise",
) -> Any:
"""
Get an object into the band storage.
Slicing support is also provided.
"""

async def get_info(
self,
session_id: str,
data_key: str,
level: StorageLevel = None,
) -> List[InternalDataInfo]:
"""
Get internal information of an object
"""

async def pin(
self,
session_id: str,
data_key: str,
level: StorageLevel = None,
error: str = "raise",
):
"""
Pin specific data on a specific level.
The object will get a read-only lock until unpinned
"""

async def unpin(
self,
session_id: str,
data_key: str,
level: StorageLevel = None,
error: str = "raise",
):
"""
Unpin specific data on a specific level
"""
```

A `StorageItemManagerActor` will hold all information necessary for kernel data management. It comprises of four separate handlers, namely `QuotaHandler`, `LockHandler`, `MetaHandler` and `ReferenceHandler`, implemented separately to reduce potential call overhead. Note that this actor only deal with data metas, not data themselves. Data are handled in caller actors with storage backends.

## User Storage

User storage API wraps kernel storage and provides more capabilities including multiple level handling, spill and transfer. The API can look like

```python
StorageLevels = Optional[List[StorageLevel]]

class UserStorageAPI:
@classmethod
async def create(
cls,
session_id: str,
band_name: str,
worker_address: str,
) -> "UserStorageAPI":
"""
Create a session and band specific API
"""

async def fetch(
self,
data_key: str,
levels: StorageLevels = None,
band_name: str = None,
remote_address: str = None,
error: str = "raise",
):
"""
Fetch object from remote worker or load object from disk
"""

async def open_reader(
self,
data_key: str,
levels: StorageLevels = None,
) -> UserStorageFileObject:
"""
Create a reader on a specific file
"""

async def open_writer(
self,
data_key: str,
size: int,
levels: StorageLevels = None,
band_name: str = None,
) -> UserStorageFileObject:
"""
Create a writer on a specific file
"""

async def delete(self, data_key: str, error: str = "raise"):
"""
Delete a file with specified keys
"""

async def put(
self,
data_key: str,
obj: Any,
levels: StorageLevels = None,
band_name: str = None,
) -> InternalDataInfo:
"""
Put an object into the band storage
"""

async def get(
self,
data_key: str,
conditions: List = None,
levels: StorageLevels = None,
band_name: str = None,
error: str = "raise",
) -> Any:
"""
Get an object into the band storage.
Slicing support is also provided.
"""

async def get_info(
self,
data_key: str,
levels: StorageLevels = None,
band_name: str = None,
) -> List[InternalDataInfo]:
"""
Get internal information of an object
"""

async def pin(
self,
data_key: str,
levels: StorageLevels = None,
band_name: str = None,
error: str = "raise",
):
"""
Pin specific data on a specific level.
The object will get a read-only lock until unpinned
"""

async def unpin(
self,
session_id: str,
data_key: str,
level: StorageLevel = None,
band_name: str = None,
error: str = "raise",
):
"""
Unpin specific data on a specific level
"""
```

## Spill

To implement spill, We need a `SpillManagerActor` to coordinate spill actions. A spill actor will be look like

```python
class SpillManagerActor(mo.StatelessActor):
@classmethod
def gen_uid(cls, band_name: str, storage_level: int) -> str:
pass

def notify_spillable(self, data_key: str, size: int):
"""
Register a spillable data key.
Only called when spill state is True.
"""

async def acquire_spill_lock(self, size: int) -> List[str]:
"""
Acquire certain size for spill and lock the actor
for spill. Keys will be returned for the caller to
spill.
"""

def release_spill_lock(self):
"""
Release the actor when spill ends.
"""

def wait_spill_state_change(self, last_state: bool) -> bool:
"""
Wait until the state of spill changes.
"""
```

Inside the actor, we define a boolean state to indicate whether the storage level is under spill. When the state changes to True, it will be broadcasted to all subscribers to notify them to notify data changes. When the storage is about to spill, it calls `acquire_spill_lock` and supply some sizes. Then the actor enters spill state, locks the actor and then checks for keys to spill. When sizes to spill is available, it will return keys to spill and spill is started from the caller. When spill ends (finishes or encounters an error), the caller calls `release_spill_lock` to release the spill lock for other callers. When there is no pending callers, the state of the actor turns into False.

## Transfer / Remote IO

To implement data transfer, we propose a two-actor solution. We will add a `SenderManagerActor` and a `RemoteIOActor` to do all required things. The `SenderManagerActor` masters data transfer initiated between workers, and `RemoteIOActor` handles remote readers and writers both for inter-worker data transfer as well as `UserStorageAPI`.

When starting an inter-worker transfer, a request is sent to `SenderManagerActor` at the worker hosting the data to send to the calling worker. It calls `RemoteIOActor.create_writer` at receiver site and then `write_data` with batch calls.

`RemoteIOActor` will look like

```python
class RemoteIOActor(mo.StatelessActor):
@mo.batch
async def create_reader(
self,
session_id: str,
data_key: str,
levels: StorageLevels,
) -> List[str]:
pass

@mo.batch
async def create_writer(
self,
session_id: str,
data_key: str,
data_size: int,
levels: StorageLevels,
) -> List[str]:
pass

@mo.batch
async def read_data(
self,
session_id: str,
reader_key: str,
data_buffer: bytes,
size: int,
):
pass

@mo.batch
async def write_data(
self,
session_id: str,
writer_key: str,
data_buffer: bytes,
is_eof: bool,
):
pass

@mo.batch
async def close(
self,
session_id: str,
key: str,
):
pass
```

And `SenderManagerActor` will look like

```python
class SenderManagerActor(mo.StatelessActor):
@mo.extensible
async def send_batch_data(
self,
session_id: str,
data_keys: List[str],
address: str,
level: StorageLevel,
band_name: str = "numa-0",
block_size: int = None,
error: str = "raise",
):
pass
```

The proposal names KernelStorageAPI, UserStorageAPI, SpillManagerActor, RemoteIOActor, and SenderManagerActor; start by locating the existing storage and actor implementations corresponding to these entry points. Compare their current responsibilities with the proposed kernel/user storage split and transfer interfaces. Done means the storage design addresses spill coordination, remote I/O, flexible locations, and reduced IPC calls as described.

Written by the indexing model from the issue text.