init commit

nimbus/components/dump/__init__.py

@@ -0,0 +1,7 @@
+from .base_dumper import BaseDumper
+
+dumper_dict = {}
+
+
+def register(type_name: str, cls: BaseDumper):
+    dumper_dict[type_name] = cls

nimbus/components/dump/base_dumper.py

@@ -0,0 +1,82 @@
+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")