83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
import time
|
|
from abc import abstractmethod
|
|
|
|
from pympler import asizeof
|
|
|
|
from nimbus.components.data.iterator import Iterator
|
|
from nimbus.components.data.package import Package
|
|
from nimbus.utils.utils import unpack_iter_data
|
|
|
|
|
|
class BaseDumper(Iterator):
|
|
def __init__(self, data_iter, output_queue, max_queue_num=1):
|
|
super().__init__()
|
|
self.data_iter = data_iter
|
|
self.scene = None
|
|
self.output_queue = output_queue
|
|
self.total_case = 0
|
|
self.success_case = 0
|
|
self.max_queue_num = max_queue_num
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def _next(self):
|
|
try:
|
|
data = next(self.data_iter)
|
|
scene, seq, obs = unpack_iter_data(data)
|
|
self.total_case += 1
|
|
if scene is not None:
|
|
if self.scene is not None and (
|
|
scene.task_id != self.scene.task_id
|
|
or scene.name != self.scene.name
|
|
or scene.task_exec_num != self.scene.task_exec_num
|
|
):
|
|
self.logger.info(
|
|
f"Scene {self.scene.name} generate finish, success rate: {self.success_case}/{self.total_case}"
|
|
)
|
|
self.total_case = 1
|
|
self.success_case = 0
|
|
self.scene = scene
|
|
if obs is None and seq is None:
|
|
self.logger.info(f"generate failed, skip once! success rate: {self.success_case}/{self.total_case}")
|
|
if self.scene is not None:
|
|
self.scene.update_generate_status(success=False)
|
|
return None
|
|
io_start_time = time.time()
|
|
if self.output_queue is not None:
|
|
obj = self.dump(seq, obs)
|
|
pack = Package(obj, task_id=scene.task_id, task_name=scene.name)
|
|
pack.serialize()
|
|
|
|
wait_time = time.time()
|
|
while self.output_queue.qsize() >= self.max_queue_num:
|
|
time.sleep(1)
|
|
end_time = time.time()
|
|
self.collect_wait_time_info(end_time - wait_time)
|
|
|
|
st = time.time()
|
|
self.output_queue.put(pack)
|
|
ed = time.time()
|
|
self.logger.info(f"put time: {ed - st}, data size: {asizeof.asizeof(obj)}")
|
|
else:
|
|
obj = self.dump(seq, obs)
|
|
self.success_case += 1
|
|
self.scene.update_generate_status(success=True)
|
|
self.collect_seq_info(1, time.time() - io_start_time)
|
|
except StopIteration:
|
|
if self.output_queue is not None:
|
|
pack = Package(None, stop_sig=True)
|
|
self.output_queue.put(pack)
|
|
if self.scene is not None:
|
|
self.logger.info(
|
|
f"Scene {self.scene.name} generate finish, success rate: {self.success_case}/{self.total_case}"
|
|
)
|
|
raise StopIteration("no data")
|
|
except Exception as e:
|
|
self.logger.exception(f"Error during data dumping: {e}")
|
|
raise e
|
|
|
|
@abstractmethod
|
|
def dump(self, seq, obs):
|
|
raise NotImplementedError("This method should be overridden by subclasses")
|