_{Comment @coderabbitai help to get the list of available commands and usage tips.}

There is a pubic utils for zero-copy serialization in #140 , maybe it can be directly used here so that all the backends can share the same set of serial/deserial utils to improve maintainability.

There is a pubic utils for zero-copy serialization in #140 , maybe it can be directly used here so that all the backends can share the same set of serial/deserial utils to improve maintainability.

Yep I noticed this utility. As I mentioned in the PR description, yuanrong kv client expects a memoryview as an argument. Is it possible to extract the memoryview from the return value of serialization?

I see. There is a little bit tricky since both MsgpackEncoder and the general serialization cannot directly return only the memory view obj. If yuanrong do not need to tackle tensordict as a whole or nested tensor (which should be already unbinded to tensors), I think it's ok to use memoryview(obj.numpy()) directly.

Need to make sure all these scenarios are considered:

  nested_tensor = torch.nested.as_nested_tensor(
      [torch.randn(4, 3), torch.randn(2, 4)], layout=torch.strided
  ),
  jagged_tensor= torch.nested.as_nested_tensor(
      [torch.randn(4, 5), torch.randn(4, 54)], layout=torch.jagged
  ),

Theoretically, these nested tensors should be splited into single tensors during making values. But in practice, I recommend add more check to make sure there is not nested tensor trying to use memory view.

tensor.is_nested and tensor.layout == torch.strided

After discussion, we agree that serial.utils.serialization can extract tensor and non-tensor, but for one key yuanrong can only handle one memoryview/bytes/bytearray/str. In other words, yuanrong mset interface is not as flexible as vanilla storage unit implementation. In fact, it's possible taht other backend has the same semantics as yuanrong. I'm afraid they can't utilize serial.utils.serialization either. So here we can accept a small zcopy helper inside yuarnong impl.

The test result shows that zcopy brings 5~20% mitigation of total latency of put, and 2x get throughput compared to clone each view tensor.

btw, where is the corresponding deserialization for memory view zero copy?

btw, where is the corresponding deserialization for memory view zero copy?

you are right, I'm redesigning deserialization to make it work for both pickled bytes and memoryview. it's expected be reused by other kv backends

I'm afraid it's impossible to achieve the following targets at the same time

1. use current set(keys, valus) where vals are memoryview or bytes or bytearray or str, and get(keys) -> list[bytes]
2. support custom zero-copy ser/des for certain types

The proof: if we want to support custom ser/des, we need metadata and type id to decide which deserializer to use on a bytes. These extra information can be accessed by either 1. concat with real value; 2. store in another kv pair. Neither of them is an ideal choice.

1. concat still causes copy. e.g. np.ndarray -> concat with metadata as one big bytes -> set
2. new kv pair doubles the kv management pressure of control plane

what we need is something like np.ndarray --(directly write into)--> rdma memory region, only one copy. It's possible if set can provide a buffer from registered rdma memory region for users to write their array and metadata into. It's a feasible way since yuanrong already supports this in c++ but not in python.

zmq can support this because of send_multipart. It transfers multiple contiguous memory at the same time. The receiver can parse some of them as metadata and some of them as real data, so that custom deserializer is possible. This is another way to tackle: set supports one key -> multiple contiguous memory (e.g. metadata and real data memoryview).

In summary, two feasible ways to achieve custom zero-copy deserializer with kv system:

1. kv system supports create_buffer + set semantics: directly write to memory region, only one copy

metadata = pickle.dumps(prepare_metadata(array))
size = len(metadata) + len(array)
buffers = mcreate([key], [size])
buffer = buffers[0]
buffer[:len(metadata)] = metadata
buffer[len(metadata)] = memoryview(array)
kv_client.mset(buffers)

2. kv system supports key : list[val] to avoid concat on user side

metadata = pickle.dumps(prepare_metadata(array))
kv_client.mset([key], [[metadata, memoryview(array)]])

I'm afraid it's impossible to achieve the following targets at the same time

1. use current set(keys, valus) where vals are memoryview or bytes or bytearray or str, and get(keys) -> list[bytes]
2. support custom zero-copy ser/des for certain types

The proof: if we want to support custom ser/des, we need metadata and type id to decide which deserializer to use on a bytes. These extra information can be accessed by either 1. concat with real value; 2. store in another kv pair. Neither of them is an ideal choice.

1. concat still causes copy. e.g. np.ndarray -> concat with metadata as one big bytes -> set
2. new kv pair doubles the kv management pressure of control plane

what we need is something like np.ndarray --(directly write into)--> rdma memory region, only one copy. It's possible if set can provide a buffer from registered rdma memory region for users to write their array and metadata into. It's a feasible way since yuanrong already supports this in c++ but not in python.

zmq can support this because of send_multipart. It transfers multiple contiguous memory at the same time. The receiver can parse some of them as metadata and some of them as real data, so that custom deserializer is possible. This is another way to tackle: set` supports one key -> multiple contiguous memory (e.g. metadata and real data memoryview).

In summary, two feasible ways to achieve custom zero-copy deserializer with kv system:

1. kv system supports create_buffer + set semantics: directly write to memory region, only one copy
2. kv system supports key : list[val] to avoid concat on user side.

Could we transfer metadata into tensor？And then add metadata to tensordata using view or others if memoryview support non-contiguous tensor.

Could we transfer metadata into tensor？And then add metadata to tensordata using view or others if memoryview support non-contiguous tensor.

so far I can't think of a good idea, becausememoryview can only support non-contiguous memory segments that are originally from one big large memory segment (e.g. memoryview(data[::2])), but memoryview(array, metadata) will fail.

1. kv system supports create_buffer + set semantics: directly write to memory region, only one copy

For this approach, how can the receiver deserialize the data? (how can we know the size of metadata part)?

I'm not sure what is the drawback of this implementation: we store the metadata and give it a dedicated key (by appending a suffix like _meta/_data to the original key), treating it as a separate object alongside the raw data.

This will essentially double the total key-value pairs in the system, and we need to ensure atomicity between the _meta and _data keys—if only one of them is successfully operated (e.g., data is updated but metadata fails to write), it will result in invalid metadata-data mappings, leading to deserialization failures or data corruption.

Is that acceptable?

metadata = pickle.dumps(prepare_metadata(array))
kv_client.put(keys=['key_meta', key_data], values=[metadata, array])

For this approach, how can the receiver deserialize the data? (how can we know the size of metadata part)?

we can either set the length of metadata as a constant, or prepend the length of metadata at the beginning. for example, metadata_length (uint32)

mset([key], [[metadata_length, metadata, array]])

a little dirty :)

Ya I'll test if meta suffix significantly affects the latency/throughput of yuanrong

I wrote a sample use case of one-copy set/get on CPU side. The following code will be written in tq and hopefully can be merged to yuanrong

def mset_zcopy(keys: list[str], objects: list[Any]):
	# prepare metadata and memoryview
	for obj in objects:
		md, mv = serialize(obj) # if any, pickle
		serialized_metas.append(pickle.dumps(md))
		memoryviews.append(mv)
		metalen = len(serialized_metas[-1]) # HEADER_SIZE = 4 bytes
		sizes.append(HEADER_SIZE + metalen + len(mv))
		# | HEADER | SERIALIZED_META | MEMORYVIEW |

	# copy metadata and memoryviews to rdma buffer
	buffers = kv_client.mcreate(keys, sizes)
	for buffer, md, mv in zip(buffers, serialized_metas, memoryviews):
		fill_header(buffer, len(md))
		buffer[HEADER_SIZE : HEADER_SIZE + len(HEADER_SIZE)] = md
		cursor = HEADER_SIZE + len(HEADER_SIZE)
		buffer[cursor : cursor + len(mv)] = mv # one copy

	kv_client.mset(buffers)

def mget_zcopy(keys: list[str]) -> list[Any]:
	values = kv_client.get(keys)
	for val in values:
		metalen = parse_header(val)
		md = pickle.loads(val[HEADER_SIZE : HEADER_SIZE + metalen])
		mv = memoryview(val[HEADER_SIZE + metalen :])
		objects.append(deserialize(md, mv))
	return objects

where serialize can utilize pickle.dumps(protocol=pickle.HIGHEST_PROTOCOL)

def serialize(obj):
	mvs = []
	meta = pickle.dumps(obj, buffer_callback=mvs.append, protocol=pickle.HIGHEST_PROTOCOL)
	assert len(mvs) == 1
	return meta, memoryview(mv[0])

For simplicity, I assume only one contiguous memoryview. To support objects with multiple contiguous chunks, some logic needs to be modified.