init commit

2026-03-16 11:44:10 +00:00
commit 94384a93c9
552 changed files with 363038 additions and 0 deletions
--- a/nimbus/init.py
+++ b/nimbus/init.py
@@ -0,0 +1,16 @@
+import ray
+
+from nimbus.utils.types import STAGE_PIPE
+
+from .data_engine import DataEngine, DistPipeDataEngine
+
+
+def run_data_engine(config, master_seed=None):
+    import nimbus_extension  # noqa: F401  pylint: disable=unused-import
+
+    if STAGE_PIPE in config:
+        ray.init(num_gpus=1)
+        data_engine = DistPipeDataEngine(config, master_seed=master_seed)
+    else:
+        data_engine = DataEngine(config, master_seed=master_seed)
+    data_engine.run()