Function ram_read_write_worker

det/projects/DensePose/tests/test_tensor_storage.py:122–180

()

def ram_read_write_worker():
    schema = {
        "tf": SizeData(dtype="float32", shape=(112, 112)),
        "ti": SizeData(dtype="int32", shape=(4, 64, 64)),
    }
    storage = SingleProcessRamTensorStorage(schema, io.BytesIO())
    world_size = comm.get_world_size()
    rank = comm.get_rank()
    data_elts = []
    # prepare different number of tensors in different processes
    for i in range(rank + 1):
        data_elt = {
            "tf": torch.ones((112, 112), dtype=torch.float32) * (rank + i * world_size),
            "ti": torch.ones((4, 64, 64), dtype=torch.int32) * (rank + i * world_size),
        }
        data_elts.append(data_elt)
    # write data to the single process storage
    for i in range(rank + 1):
        record_id = storage.put(data_elts[i])
        assert record_id == i, f"Process {rank}: record ID {record_id}, expected {i}"
    comm.synchronize()
    # gather all data in process rank 0
    multi_storage = storage_gather(storage)
    if rank != 0:
        return
    # read and check data from the multiprocess storage
    for j in range(world_size):
        for i in range(j):
            record = multi_storage.get(j, i)
            record_gt = {
                "tf": torch.ones((112, 112), dtype=torch.float32) * (j + i * world_size),
                "ti": torch.ones((4, 64, 64), dtype=torch.int32) * (j + i * world_size),
            }
            assert len(record) == len(schema), (
                f"Process {rank}: multi storage record, rank {j}, id {i}: "
                f"expected {len(schema)} fields in the record, got {len(record)}"
            )
            for field_name in schema:
                assert field_name in record, (
                    f"Process {rank}: multi storage record, rank {j}, id {i}: "
                    f"field {field_name} not in the record"
                )

                assert record_gt[field_name].shape == record[field_name].shape, (
                    f"Process {rank}: multi storage record, rank {j}, id {i}: "
                    f"field {field_name}, expected shape {record_gt[field_name].shape} "
                    f"got {record[field_name].shape}"
                )
                assert record_gt[field_name].dtype == record[field_name].dtype, (
                    f"Process {rank}: multi storage record, rank {j}, id {i}: "
                    f"field {field_name}, expected dtype {record_gt[field_name].dtype} "
                    f"got {record[field_name].dtype}"
                )
                assert torch.allclose(record_gt[field_name], record[field_name]), (
                    f"Process {rank}: multi storage record, rank {j}, id {i}: "
                    f"field {field_name}, tensors are not close enough:"
                    f"L-inf {(record_gt[field_name]-record[field_name]).abs_().max()} "
                    f"L1 {(record_gt[field_name]-record[field_name]).abs_().sum()} "