compare read/write consistency in e2e tests

Author: mag1c-h

ucm/store/test/e2e/cache_on_posix_test.py

@@ -28,7 +28,6 @@
 from abc import ABC
 from typing import List
 
-import cupy
 import torch
 
 from ucm.store.cache.connector import UcmCacheStore
@@ -43,6 +42,7 @@ def __init__(
         layer_size: int,
         chunk_size: int,
         storage_backends: List[str],
+        device_id: int,
     ):
         super().__init__()
         chunk_block_size = tensor_size * layer_size * chunk_size
@@ -57,7 +57,7 @@ def __init__(
         cache_config = {}
         cache_config["backend"] = self.posix.cc_store()
         cache_config["engine_id"] = secrets.token_hex(8)
-        cache_config["device_id"] = 1
+        cache_config["device_id"] = device_id
         cache_config["tensor_size"] = tensor_size
         cache_config["shard_size"] = chunk_block_size
         cache_config["block_size"] = chunk_block_size
@@ -111,38 +111,64 @@ def check(self, task: Task) -> bool:
         return self.cache.check(task)
 
 
-def main():
-    tensor_size = 262144
-    layer_size = 64
-    chunk_size = 4
-    request_size = chunk_size * 16
-    storage_backends = ["."]
-    store = HierarchicalStore(tensor_size, layer_size, chunk_size, storage_backends)
+def cmp_and_print_diff(a, b, rtol=0.0, atol=0.0):
+    for r, (row_a, row_b) in enumerate(zip(a, b)):
+        for c, (ta, tb) in enumerate(zip(row_a, row_b)):
+            if not torch.allclose(ta, tb, rtol=rtol, atol=atol):
+                mask = ~torch.isclose(ta, tb, rtol=rtol, atol=atol)
+                diff_a = ta[mask].cpu()
+                diff_b = tb[mask].cpu()
+                print(f"DIFF at [{r}][{c}]  total {mask.sum().item()} element(s)")
+                print("  a val:", diff_a.flatten())
+                print("  b val:", diff_b.flatten())
+                assert False
+
+
+def e2e_test(
+    store: HierarchicalStore,
+    tensor_size: int,
+    layer_size: int,
+    chunk_size: int,
+    request_size: int,
+    device_id: int,
+):
     chunk_block_ids = [secrets.token_bytes(16) for _ in range(request_size)]
     founds = store.lookup(chunk_block_ids)
     assert not all(founds)
     shard_indexes = [0 for _ in range(request_size)]
-    src_addrs = [
+    src_tensors = [
         [
-            cupy.cuda.alloc_pinned_memory(tensor_size).ptr
+            torch.rand(
+                [tensor_size // 2],
+                dtype=torch.bfloat16,
+                device="cuda:{}".format(device_id),
+            )
             for _ in range(layer_size * chunk_size)
         ]
         for _ in range(request_size)
     ]
-    task = store.dump_data(chunk_block_ids, shard_indexes, src_addrs)
+    task = store.dump(chunk_block_ids, shard_indexes, src_tensors)
     store.wait(task)
     time.sleep(1)
-    dst_addrs = [
-        [
-            cupy.cuda.alloc_pinned_memory(tensor_size).ptr
-            for _ in range(layer_size * chunk_size)
-        ]
-        for _ in range(request_size)
-    ]
-    founds = store.lookup(chunk_block_ids)
-    assert all(founds)
-    task = store.load_data(chunk_block_ids, shard_indexes, dst_addrs)
+    dst_tensors = [[torch.empty_like(t) for t in row] for row in src_tensors]
+    task = store.load(chunk_block_ids, shard_indexes, dst_tensors)
     store.wait(task)
+    cmp_and_print_diff(src_tensors, dst_tensors)
+
+
+def main():
+    tensor_size = 262144
+    layer_size = 64
+    chunk_size = 4
+    request_size = chunk_size * 16
+    storage_backends = ["."]
+    device_id = 1
+    test_batch_number = 64
+    store = HierarchicalStore(
+        tensor_size, layer_size, chunk_size, storage_backends, device_id
+    )
+    for _ in range(test_batch_number):
+        e2e_test(store, tensor_size, layer_size, chunk_size, request_size, device_id)
 
 
 if __name__ == "__main__":