Main process manages models, child process loads models.

Author: RkGrit

iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/basic_pipeline.py

@@ -21,14 +21,14 @@
 import torch
 
 from iotdb.ainode.core.exception import InferenceModelInternalError
-from iotdb.ainode.core.manager.model_manager import ModelManager
+from iotdb.ainode.core.model.model_loader import load_model
 
 
 class BasicPipeline(ABC):
- def __init__(self, model_id, **infer_kwargs):
- self.model_id = model_id
+ def __init__(self, model_info, **infer_kwargs):
+ self.model_info = model_info
  self.device = infer_kwargs.get("device", "cpu")
- self.model = ModelManager().load_model(model_id, device_map=self.device)
+ self.model = load_model(model_info, device_map=self.device)
 
  def _preprocess(self, inputs):
  """
@@ -45,8 +45,8 @@ def _postprocess(self, output: torch.Tensor):
 
 
 class ForecastPipeline(BasicPipeline):
- def __init__(self, model_id, **infer_kwargs):
- super().__init__(model_id, infer_kwargs=infer_kwargs)
+ def __init__(self, model_info, **infer_kwargs):
+ super().__init__(model_info, infer_kwargs=infer_kwargs)
 
  def _preprocess(self, inputs):
  if len(inputs.shape) != 2:
@@ -63,8 +63,8 @@ def _postprocess(self, output: torch.Tensor):
 
 
 class ClassificationPipeline(BasicPipeline):
- def __init__(self, model_id, **infer_kwargs):
- super().__init__(model_id, infer_kwargs=infer_kwargs)
+ def __init__(self, model_info, **infer_kwargs):
+ super().__init__(model_info, infer_kwargs=infer_kwargs)
 
  def _preprocess(self, inputs):
  pass
@@ -80,8 +80,8 @@ def _postprocess(self, output: torch.Tensor):
 
 
 class ChatPipeline(BasicPipeline):
- def __init__(self, model_id, **infer_kwargs):
- super().__init__(model_id, infer_kwargs=infer_kwargs)
+ def __init__(self, model_info, **infer_kwargs):
+ super().__init__(model_info, infer_kwargs=infer_kwargs)
 
  def _preprocess(self, inputs):
  pass

iotdb-core/ainode/iotdb/ainode/core/inference/pipeline/pipeline_loader.py

@@ -49,4 +49,4 @@ def load_pipeline(model_info: ModelInfo, device: str, **kwargs):
  model_info.model_id, model_info.pipeline_cls
  )
 
- return pipeline_cls(model_info.model_id, device=device)
+ return pipeline_cls(model_info, device=device)

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

@@ -21,7 +21,7 @@
 from iotdb.ainode.core.constant import TSStatusCode
 from iotdb.ainode.core.exception import BuiltInModelDeletionError
 from iotdb.ainode.core.log import Logger
-from iotdb.ainode.core.model.model_loader import ModelLoader
+from iotdb.ainode.core.model.model_loader import load_model
 from iotdb.ainode.core.model.model_storage import ModelCategory, ModelInfo, ModelStorage
 from iotdb.ainode.core.rpc.status import get_status
 from iotdb.ainode.core.util.decorator import singleton
@@ -41,7 +41,6 @@
 class ModelManager:
  def __init__(self):
  self._model_storage = ModelStorage()
- self._model_loader = ModelLoader(storage=self._model_storage)
 
  def register_model(
  self,
@@ -75,7 +74,8 @@ def delete_model(self, req: TDeleteModelReq) -> TSStatus:
  return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
 
  def load_model(self, model_id: str, **kwargs) -> Any:
- return self._model_loader.load_model(model_id=model_id, **kwargs)
+ model_info = self.get_model_info(model_id)
+ return load_model(model_info=model_info, **kwargs)
 
  def get_model_info(
  self,

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

@@ -25,6 +25,7 @@
 from iotdb.ainode.core.exception import ModelNotExistError
 from iotdb.ainode.core.log import Logger
 from iotdb.ainode.core.manager.model_manager import ModelManager
+from iotdb.ainode.core.model.model_loader import load_model
 
 logger = Logger()
 
@@ -46,7 +47,8 @@ def measure_model_memory(device: torch.device, model_id: str) -> int:
  torch.cuda.synchronize(device)
  start = torch.cuda.memory_reserved(device)
 
- model = ModelManager().load_model(model_id).to(device)
+ model_info = ModelManager().get_model_info(model_id)
+ model = load_model(model_info).to(device)
  torch.cuda.synchronize(device)
  end = torch.cuda.memory_reserved(device)
  usage = end - start

iotdb-core/ainode/iotdb/ainode/core/model/model_info.py

@@ -28,6 +28,7 @@ def __init__(
  category: ModelCategory,
  state: ModelStates,
  model_type: str = "",
+ config_cls: str = "",
  model_cls: str = "",
  pipeline_cls: str = "",
  repo_id: str = "",
@@ -38,6 +39,7 @@ def __init__(
  self.model_type = model_type
  self.category = category
  self.state = state
+ self.config_cls = config_cls
  self.model_cls = model_cls
  self.pipeline_cls = pipeline_cls
  self.repo_id = repo_id
@@ -112,6 +114,7 @@ def __repr__(self):
  category=ModelCategory.BUILTIN,
  state=ModelStates.INACTIVE,
  model_type="timer",
+ config_cls="configuration_timer.TimerConfig",
  model_cls="modeling_timer.TimerForPrediction",
  pipeline_cls="pipeline_timer.TimerPipeline",
  repo_id="thuml/timer-base-84m",
@@ -121,6 +124,7 @@ def __repr__(self):
  category=ModelCategory.BUILTIN,
  state=ModelStates.INACTIVE,
  model_type="sundial",
+ config_cls="configuration_sundial.SundialConfig",
  model_cls="modeling_sundial.SundialForPrediction",
  pipeline_cls="pipeline_sundial.SundialPipeline",
  repo_id="thuml/sundial-base-128m",