[AINode] Limit max inference length (apache#15982)

Author: CRZbulabula

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

@@ -134,7 +134,7 @@ def start(self):
                 raise e
 
         # Start the RPC service
-        self._rpc_handler = AINodeRPCServiceHandler(aiNode=self)
+        self._rpc_handler = AINodeRPCServiceHandler(ainode=self)
         self._rpc_service = AINodeRPCService(self._rpc_handler)
         self._rpc_service.start()
         self._rpc_service.join(1)

iotdb-core/ainode/ainode/core/config.py

@@ -31,6 +31,7 @@
     AINODE_CONF_GIT_FILE_NAME,
     AINODE_CONF_POM_FILE_NAME,
     AINODE_INFERENCE_BATCH_INTERVAL_IN_MS,
+    AINODE_INFERENCE_MAX_PREDICT_LENGTH,
     AINODE_LOG_DIR,
     AINODE_MODELS_DIR,
     AINODE_ROOT_CONF_DIRECTORY_NAME,
@@ -72,6 +73,9 @@ def __init__(self):
         self._ain_inference_batch_interval_in_ms: int = (
             AINODE_INFERENCE_BATCH_INTERVAL_IN_MS
         )
+        self._ain_inference_max_predict_length: int = (
+            AINODE_INFERENCE_MAX_PREDICT_LENGTH
+        )
 
         # log directory
         self._ain_logs_dir: str = AINODE_LOG_DIR
@@ -140,6 +144,14 @@ def set_ain_inference_batch_interval_in_ms(
     ) -> None:
         self._ain_inference_batch_interval_in_ms = ain_inference_batch_interval_in_ms
 
+    def get_ain_inference_max_predict_length(self) -> int:
+        return self._ain_inference_max_predict_length
+
+    def set_ain_inference_max_predict_length(
+        self, ain_inference_max_predict_length: int
+    ) -> None:
+        self._ain_inference_max_predict_length = ain_inference_max_predict_length
+
     def get_ain_logs_dir(self) -> str:
         return self._ain_logs_dir
 

iotdb-core/ainode/ainode/core/constant.py

@@ -50,6 +50,7 @@
 
 # AINode inference configuration
 AINODE_INFERENCE_BATCH_INTERVAL_IN_MS = 15
+AINODE_INFERENCE_MAX_PREDICT_LENGTH = 2880
 
 # AINode folder structure
 AINODE_ROOT_DIR = os.path.dirname(

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

@@ -23,7 +23,7 @@
 import numpy as np
 import torch
 import torch.multiprocessing as mp
-from transformers import PretrainedConfig, PreTrainedModel
+from transformers import PretrainedConfig
 
 from ainode.core.config import AINodeDescriptor
 from ainode.core.inference.inference_request import InferenceRequest
@@ -46,7 +46,7 @@ class InferenceRequestPool(mp.Process):
     def __init__(
         self,
         pool_id: int,
-        model_id: int,
+        model_id: str,
         config: PretrainedConfig,
         request_queue: mp.Queue,
         result_queue: mp.Queue,
@@ -58,6 +58,7 @@ def __init__(
         self.config = config
         self.pool_kwargs = pool_kwargs
         self.model = None
+        self._model_manager = None
         self.device = None
 
         # TODO: A scheduler is necessary for better handling following queues

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

@@ -22,7 +22,7 @@
 from transformers.modeling_outputs import MoeCausalLMOutputWithPast
 
 
-def _generate_req_id(length=10, charset=string.ascii_letters + string.digits) -> str:
+def generate_req_id(length=10, charset=string.ascii_letters + string.digits) -> str:
     """
     Generate a random req_id string of specified length.
     The length is 10 by default, with 10^{17} possible combinations.

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

@@ -30,6 +30,7 @@
 from ainode.core.exception import (
     InferenceModelInternalError,
     InvalidWindowArgumentError,
+    NumericalRangeException,
     runtime_error_extractor,
 )
 from ainode.core.inference.inference_request import (
@@ -40,7 +41,7 @@
 from ainode.core.inference.strategy.timer_sundial_inference_pipeline import (
     TimerSundialInferencePipeline,
 )
-from ainode.core.inference.utils import _generate_req_id
+from ainode.core.inference.utils import generate_req_id
 from ainode.core.log import Logger
 from ainode.core.manager.model_manager import ModelManager
 from ainode.core.model.sundial.configuration_sundial import SundialConfig
@@ -214,6 +215,20 @@ def _run(
             full_data = deserializer(raw)
             inference_attrs = extract_attrs(req)
 
+            predict_length = inference_attrs.get("predict_length", 96)
+            if (
+                predict_length
+                > AINodeDescriptor().get_config().get_ain_inference_max_predict_length()
+            ):
+                raise NumericalRangeException(
+                    "output_length",
+                    1,
+                    AINodeDescriptor()
+                    .get_config()
+                    .get_ain_inference_max_predict_length(),
+                    predict_length,
+                )
+
             if model_id == self.ACCELERATE_MODEL_ID and self.DEFAULT_POOL_SIZE > 0:
                 # TODO: Logic in this branch shall handle all LTSM inferences
                 # TODO: TSBlock -> Tensor codes should be unified
@@ -223,10 +238,10 @@ def _run(
                 # 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(),
+                    req_id=generate_req_id(),
                     inputs=inputs,
                     inference_pipeline=TimerSundialInferencePipeline(SundialConfig()),
-                    max_new_tokens=inference_attrs.get("predict_length", 96),
+                    max_new_tokens=predict_length,
                 )
                 infer_proxy = InferenceRequestProxy(infer_req.req_id)
                 with self._result_wrapper_lock:

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

@@ -41,13 +41,13 @@
 
 
 class AINodeRPCServiceHandler(IAINodeRPCService.Iface):
-    def __init__(self, aiNode):
-        self._aiNode = aiNode
+    def __init__(self, ainode):
+        self._ainode = ainode
         self._model_manager = ModelManager()
         self._inference_manager = InferenceManager()
 
     def stopAINode(self) -> TSStatus:
-        self._aiNode.stop()
+        self._ainode.stop()
         return get_status(TSStatusCode.SUCCESS_STATUS, "AINode stopped successfully.")
 
     def registerModel(self, req: TRegisterModelReq) -> TRegisterModelResp: