File size: 6,092 Bytes
079c32c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import os
import timeit
import pytest
import tempfile
import shutil
import numpy as np
import torch
import treetensor.torch as ttorch
from ding.data.shm_buffer import ShmBuffer
from ding.data.storage_loader import FileStorageLoader
from time import sleep, time
from os import path
from ding.framework.supervisor import RecvPayload
@pytest.mark.tmp # gitlab ci and local test pass, github always fail
def test_file_storage_loader():
tempdir = path.join(tempfile.gettempdir(), "test_storage_loader")
loader = FileStorageLoader(dirname=tempdir)
try:
total_num = 200
storages = []
for i in range(10):
# 21MB
data = [
{
"s": "abc",
"obs": np.random.rand(4, 84, 84).astype(np.float32),
# "next_obs": np.random.rand(4, 84, 84).astype(np.float32),
# "obs": torch.rand(4, 84, 84, dtype=torch.float32),
"next_obs": torch.rand(4, 84, 84, dtype=torch.float32)
} for _ in range(96)
]
storage = loader.save(data)
storages.append(storage)
start = time()
for i in range(total_num):
storage = storages[i % 10]
data = storage.load()
origin_time_cost = time() - start
print("Load time cost: {:.4f}s".format(origin_time_cost))
call_times = 0
def callback(data):
assert data[0]['obs'] is not None
nonlocal call_times
call_times += 1
# First initialize shared memory is very slow, discard this time cost.
start = time()
loader._first_meet(storage=storages[0], callback=callback)
print("Initialize shared memory time: {:.4f}s".format(time() - start))
start = time()
for i in range(1, total_num):
storage = storages[i % 10]
loader.load(storage, callback)
while True:
if call_times == total_num:
break
sleep(0.01)
new_time_cost = time() - start
print("Loader time cost: {:.4f}s".format(new_time_cost))
assert new_time_cost < origin_time_cost
finally:
if path.exists(tempdir):
shutil.rmtree(tempdir)
loader.shutdown()
@pytest.mark.unittest
def test_file_storage_loader_cleanup():
tempdir = path.join(tempfile.gettempdir(), "test_storage_loader")
loader = FileStorageLoader(dirname=tempdir, ttl=1)
try:
storages = []
for _ in range(4):
data = np.random.rand(4, 84, 84).astype(np.float32)
storage = loader.save(data)
storages.append(storage)
sleep(0.5)
assert len(os.listdir(tempdir)) < 4
finally:
if path.exists(tempdir):
shutil.rmtree(tempdir)
loader.shutdown()
@pytest.mark.unittest
def test_shared_object():
loader = FileStorageLoader(dirname="")
# ========== Test array ==========
obj = [{"obs": np.random.rand(100, 100)} for _ in range(10)]
shm_obj = loader._create_shm_buffer(obj)
assert len(shm_obj.buf) == len(obj) * 2
assert isinstance(shm_obj.buf[0]["obs"], ShmBuffer)
# Callback
payload = RecvPayload(proc_id=0, data=obj)
loader._shm_callback(payload=payload, shm_obj=shm_obj)
assert len(payload.data) == 10
assert [d["obs"] is None for d in payload.data]
# ========== Putback ==========
loader._shm_putback(payload=payload, shm_obj=shm_obj)
obj = payload.data
assert len(obj) == 10
for o in obj:
assert isinstance(o["obs"], np.ndarray)
assert o["obs"].shape == (100, 100)
# ========== Test dict ==========
obj = {"obs": torch.rand(100, 100, dtype=torch.float32)}
shm_obj = loader._create_shm_buffer(obj)
assert isinstance(shm_obj.buf["obs"], ShmBuffer)
payload = RecvPayload(proc_id=0, data=obj)
loader._shm_callback(payload=payload, shm_obj=shm_obj)
assert payload.data["obs"] is None
loader._shm_putback(payload=payload, shm_obj=shm_obj)
assert isinstance(payload.data["obs"], torch.Tensor)
assert payload.data["obs"].shape == (100, 100)
# ========== Test treetensor ==========
obj = {"trajectories": [ttorch.as_tensor({"obs": torch.rand(10, 10, dtype=torch.float32)}) for _ in range(10)]}
shm_obj = loader._create_shm_buffer(obj)
payload = RecvPayload(proc_id=0, data=obj)
loader._shm_callback(payload=payload, shm_obj=shm_obj)
assert len(payload.data["trajectories"]) == 10
for traj in payload.data["trajectories"]:
assert traj["obs"] is None
loader._shm_putback(payload=payload, shm_obj=shm_obj)
for traj in payload.data["trajectories"]:
assert isinstance(traj["obs"], torch.Tensor)
assert traj["obs"].shape == (10, 10)
@pytest.mark.benchmark
def test_shared_object_benchmark():
loader = FileStorageLoader(dirname="")
# ========== Test treetensor ==========
obj = {
"env_step": 0,
"trajectories": [
ttorch.as_tensor(
{
"done": False,
"reward": torch.tensor([1, 0], dtype=torch.int32),
"obs": torch.rand(4, 84, 84, dtype=torch.float32),
"next_obs": torch.rand(4, 84, 84, dtype=torch.float32),
"action": torch.tensor([1], dtype=torch.int32),
"collect_train_iter": torch.tensor([1], dtype=torch.int32),
"env_data_id": torch.tensor([1], dtype=torch.int32),
}
) for _ in range(10)
]
}
buf = loader._create_shm_buffer(obj)
payload = RecvPayload(proc_id=0, data=obj)
loader._shm_callback(payload=payload, shm_obj=buf)
def stmt():
payload.extra = buf.id_.get()
loader._shm_putback(payload=payload, shm_obj=buf)
res = timeit.repeat(stmt, repeat=5, number=1000)
print("Mean: {:.4f}s, STD: {:.4f}s, Mean each call: {:.4f}ms".format(np.mean(res), np.std(res), np.mean(res)))
assert np.mean(res) < 1
|