+ init ray in the same namespace for data processor

+ release output buffer after the active iterator is finished

Author: HYLcool

trinity/data/controllers/active_iterator.py

@@ -221,6 +221,12 @@ def run(self, thread_event: threading.Event = None):
                 traceback.print_exc()
                 return 10, "Exporting result to output buffer failed."
 
+        try:
+            dataset.release_output_buffer()
+        except Exception:
+            traceback.print_exc()
+            return -1, "Releasing output buffer failed."
+
         return 0, "success"
 
     def _group_scores(self, dataset: RftDataset) -> RftDataset:

trinity/data/core/dataset.py

@@ -49,12 +49,16 @@ def __init__(
     ):
         self.config = data_pipeline_config
         self.buffer_config = buffer_config
+        # init input buffers
         input_buffer_configs = self.config.input_buffers
         if len(input_buffer_configs) == 0:
             raise ValueError("input_buffers is empty in data pipeline config")
-        self.buffers = []
+        self.input_buffers = []
         for input_buffer_config in input_buffer_configs:
-            self.buffers.append(get_buffer_reader(input_buffer_config, self.buffer_config))
+            self.input_buffers.append(get_buffer_reader(input_buffer_config, self.buffer_config))
+        # init output buffer
+        self.output_buffer = get_buffer_writer(self.config.output_buffer, self.buffer_config)
+
         self.data = Dataset.from_list([])
         self.original_dataclass = None
 
@@ -79,28 +83,23 @@ def sort_by(self, key: str, reverse: bool = False, top_k: int = -1):
 
     def read_from_buffer(self):
         datasets = []
-        for buffer in self.buffers:
+        for buffer in self.input_buffers:
             exp_list = buffer.read()
             if self.original_dataclass is None:
                 self.original_dataclass = exp_list[0].__class__
             datasets.append(Dataset.from_list([asdict(exp) for exp in exp_list]))
         self.data = concatenate_datasets(datasets)
         logger.info(f"Read {len(self.data)} samples from input buffers")
 
-    def write_to_buffer(
-        self, output_storage_config: StorageConfig = None, buffer_config: BufferConfig = None
-    ):
-        if output_storage_config is None:
-            output_storage_config = self.config.output_buffer
-        if buffer_config is None:
-            buffer_config = self.buffer_config
-        output_buffer = get_buffer_writer(output_storage_config, buffer_config)
+    def write_to_buffer(self):
         exp_list = [dict_to_dataclass(self.original_dataclass, d) for d in self.data.to_list()]
-        output_buffer.write(exp_list)
-        output_buffer.release()
+        self.output_buffer.write(exp_list)
         logger.info(f"Wrote {len(self.data)} samples to output buffer")
         self.data = Dataset.from_list([])
 
+    def release_output_buffer(self):
+        self.output_buffer.release()
+
     def to_parquet(self, path: str):
         self.data.to_parquet(path)
 

trinity/data/server.py

@@ -1,5 +1,6 @@
 import fire
 import threading
+import ray
 from flask import Flask, jsonify, request
 from markupsafe import escape
 from typing import List
@@ -20,6 +21,9 @@ def data_processor(pipeline_type):
     config = load_config(config_path)
     config.check_and_update()
 
+    # init ray
+    ray.init(namespace=config.ray_namespace)
+
     pipeline_config = getattr(config.data_processor, pipeline_type)
     if pipeline_config is None:
         return jsonify(