[AINode] Accelerate memory efficiency of the multiprocessing architecture of inference_manager (#15956)

Author: yunbow30944

iotdb-core/ainode/ainode/core/inference/inference_request.py

@@ -59,9 +59,8 @@ def __init__(
         self.cur_step_idx = 0  # Current write position in the output step index
 
         # Preallocate output buffer [batch_size, max_new_tokens]
-        device = inputs.device
         self.output_tensor = torch.zeros(
-            self.batch_size, max_new_tokens, device=device
+            self.batch_size, max_new_tokens, device="cpu"
         )  # shape: [self.batch_size, max_new_steps]
 
     def mark_running(self):

iotdb-core/ainode/ainode/core/inference/inference_request_pool.py

@@ -28,6 +28,7 @@
 from ainode.core.config import AINodeDescriptor
 from ainode.core.inference.inference_request import InferenceRequest
 from ainode.core.log import Logger
+from ainode.core.manager.model_manager import ModelManager
 
 logger = Logger()
 
@@ -45,18 +46,19 @@ class InferenceRequestPool(mp.Process):
     def __init__(
         self,
         pool_id: int,
-        model: PreTrainedModel,
+        model_id: int,
         config: PretrainedConfig,
         request_queue: mp.Queue,
         result_queue: mp.Queue,
         **pool_kwargs,
     ):
         super().__init__()
         self.pool_id = pool_id
-        self.model = model
-        self.device = self.model.device
+        self.model_id = model_id
         self.config = config
         self.pool_kwargs = pool_kwargs
+        self.model = None
+        self.device = None
 
         # TODO: A scheduler is necessary for better handling following queues
         self._threads = []
@@ -97,19 +99,25 @@ def _step(self):
         # TODO: We need a batcher to accelerate the concurrent inference
         # TODO: Check memory size before executing requests
         request: InferenceRequest = self._running_queue.get()
+        inputs = request.inputs.to(self.device)
         output = self.model.generate(
-            request.inputs,
+            inputs,
             max_new_tokens=request.max_new_tokens,
             num_samples=10,
             revin=True,
         )
+        request.output_tensor = request.output_tensor.to(
+            self.device
+        )  # Ensure output tensor is on the same device
         request.write_step_output(output[0].mean(dim=0))
         request.inference_pipeline.post_decode()
         if request.is_finished():
             request.inference_pipeline.post_inference()
             logger.debug(
                 f"[Inference][Device-{self.device}][Pool-{self.pool_id}][ID-{request.req_id}] Request is finished"
             )
+            # ensure the output tensor is on CPU before sending to result queue
+            request.output_tensor = request.output_tensor.cpu()
             self._finished_queue.put(request)
         else:
             logger.debug(
@@ -123,6 +131,10 @@ def _requests_execute_loop(self):
             self._step()
 
     def run(self):
+        self._model_manager = ModelManager()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = self._model_manager.load_model(self.model_id, {}).to(self.device)
+
         activate_daemon = threading.Thread(
             target=self._requests_activate_loop, daemon=True
         )

iotdb-core/ainode/ainode/core/manager/inference_manager.py

@@ -135,17 +135,17 @@ def infer(self, full_data, window_interval=None, window_step=None, **_):
 
 class InferenceManager:
     ACCELERATE_MODEL_ID = "sundial"
-    DEFAULT_DEVICE = "cpu"
-    # DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # DEFAULT_DEVICE = "cpu"
+    DEFAULT_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     DEFAULT_POOL_SIZE = (
         0  # TODO: Remove these parameter by sampling model inference consumption
     )
     WAITING_INTERVAL_IN_MS = (
         AINodeDescriptor().get_config().get_ain_inference_batch_interval_in_ms()
     )  # How often to check for requests in the result queue
 
-    def __init__(self, model_manager: ModelManager):
-        self._model_manager = model_manager
+    def __init__(self):
+        self._model_manager = ModelManager()
         self._result_queue = mp.Queue()
         self._result_wrapper_map = {}
         self._result_wrapper_lock = threading.RLock()
@@ -165,14 +165,11 @@ def _init_inference_request_pool(self):
         """
         self._request_pool_map[self.ACCELERATE_MODEL_ID] = []
         for idx in range(self.DEFAULT_POOL_SIZE):
-            sundial_model = self._model_manager.load_model(
-                self.ACCELERATE_MODEL_ID, {}
-            ).to(self.DEFAULT_DEVICE)
             sundial_config = SundialConfig()
             request_queue = mp.Queue()
             request_pool = InferenceRequestPool(
                 pool_id=idx,
-                model=sundial_model,
+                model_id=self.ACCELERATE_MODEL_ID,
                 config=sundial_config,
                 request_queue=request_queue,
                 result_queue=self._result_queue,
@@ -223,7 +220,8 @@ def _run(
                 data = full_data[1][0]
                 if data.dtype.byteorder not in ("=", "|"):
                     data = data.byteswap().newbyteorder()
-                inputs = torch.tensor(data).unsqueeze(0).float().to(self.DEFAULT_DEVICE)
+                # the inputs should be on CPU before passing to the inference request
+                inputs = torch.tensor(data).unsqueeze(0).float().to("cpu")
                 infer_req = InferenceRequest(
                     req_id=_generate_req_id(),
                     inputs=inputs,

iotdb-core/ainode/ainode/core/manager/model_manager.py

@@ -29,6 +29,7 @@
 from ainode.core.model.model_info import BuiltInModelType, ModelInfo, ModelStates
 from ainode.core.model.model_storage import ModelStorage
 from ainode.core.rpc.status import get_status
+from ainode.core.util.decorator import singleton
 from ainode.thrift.ainode.ttypes import (
     TDeleteModelReq,
     TRegisterModelReq,
@@ -41,6 +42,7 @@
 logger = Logger()
 
 
+@singleton
 class ModelManager:
     def __init__(self):
         self.model_storage = ModelStorage()

iotdb-core/ainode/ainode/core/rpc/handler.py

@@ -44,7 +44,7 @@ class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
     def __init__(self, aiNode):
         self._aiNode = aiNode
         self._model_manager = ModelManager()
-        self._inference_manager = InferenceManager(model_manager=self._model_manager)
+        self._inference_manager = InferenceManager()
 
     def stopAINode(self) -> TSStatus:
         self._aiNode.stop()

iotdb-core/ainode/ainode/core/script.py

@@ -19,6 +19,8 @@
 import shutil
 import sys
 
+import torch.multiprocessing as mp
+
 from ainode.core.ainode import AINode
 from ainode.core.config import AINodeDescriptor
 from ainode.core.constant import TSStatusCode
@@ -86,6 +88,8 @@ def main():
     command = arguments[1]
     if command == "start":
         try:
+            mp.set_start_method("spawn", force=True)
+            logger.info(f"Current multiprocess start method: {mp.get_start_method()}")
             logger.info("IoTDB-AINode is starting...")
             ai_node = AINode()
             ai_node.start()