init commit

nimbus/components/load/base_randomizer.py

@@ -0,0 +1,72 @@
+import sys
+import time
+from abc import abstractmethod
+from typing import Optional
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.scene import Scene
+from nimbus.daemon.decorators import status_monitor
+
+
+class LayoutRandomizer(Iterator):
+    """
+    Base class for layout randomization in a scene. This class defines the structure for randomizing scenes and
+    tracking the randomization process. It manages the current scene, randomization count, and provides hooks for
+    subclasses to implement specific randomization logic.
+
+    Args:
+        scene_iter (Iterator): An iterator that provides scenes to be randomized.
+        random_num (int): The number of randomizations to perform for each scene before moving to the next one.
+        strict_mode (bool): If True, the randomizer will check the generation status of the current scene and retry
+                            randomization if it was not successful. This ensures that only successfully generated
+                            scenes are counted towards the randomization limit.
+    """
+
+    def __init__(self, scene_iter: Iterator, random_num: int, strict_mode: bool = False):
+        super().__init__()
+        self.scene_iter = scene_iter
+        self.random_num = random_num
+        self.strict_mode = strict_mode
+        self.cur_index = sys.maxsize
+        self.scene: Optional[Scene] = None
+
+    def reset(self, scene):
+        self.cur_index = 0
+        self.scene = scene
+
+    def _fetch_next_scene(self):
+        scene = next(self.scene_iter)
+        self.reset(scene)
+
+    @status_monitor()
+    def _randomize_with_status(self, scene) -> Scene:
+        scene = self.randomize_scene(self.scene)
+        return scene
+
+    def _next(self) -> Scene:
+        try:
+            if self.strict_mode and self.scene is not None:
+                if not self.scene.get_generate_status():
+                    self.logger.info("strict_mode is open, retry the randomization to generate sequence.")
+                    st = time.time()
+                    scene = self._randomize_with_status(self.scene)
+                    self.collect_seq_info(1, time.time() - st)
+                    return scene
+            if self.cur_index >= self.random_num:
+                self._fetch_next_scene()
+            if self.cur_index < self.random_num:
+                st = time.time()
+                scene = self._randomize_with_status(self.scene)
+                self.collect_seq_info(1, time.time() - st)
+                self.cur_index += 1
+            return scene
+        except StopIteration:
+            raise StopIteration("No more scenes to randomize.")
+        except Exception as e:
+            self.logger.exception(f"Error during scene idx {self.cur_index} randomization: {e}")
+            self.cur_index += 1
+            raise e
+
+    @abstractmethod
+    def randomize_scene(self, scene) -> Scene:
+        raise NotImplementedError("This method should be overridden by subclasses")