This is an automated email from the ASF dual-hosted git repository.
yongzao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 2b60aa5eecc [AINode] Unify .safetensors model loading through
huggingface interfaces (#15756)
2b60aa5eecc is described below
commit 2b60aa5eeccbc17a3ac5b3f8bde04d1b9455e8f4
Author: Zeyu Zhang <[email protected]>
AuthorDate: Sun Jun 22 06:48:58 2025 +0800
[AINode] Unify .safetensors model loading through huggingface interfaces
(#15756)
---
iotdb-core/ainode/ainode/core/config.py | 10 +-
iotdb-core/ainode/ainode/core/constant.py | 22 ++++-
.../ainode/core/manager/inference_manager.py | 22 ++---
.../ainode/ainode/core/manager/model_manager.py | 48 +++++++---
.../ainode/core/model/built_in_model_factory.py | 67 +++++++++-----
.../ainode/ainode/core/model/model_storage.py | 103 +++++++++++++--------
.../ainode/core/model/sundial/modeling_sundial.py | 22 +----
.../ainode/core/model/timerxl/modeling_timer.py | 22 +----
8 files changed, 184 insertions(+), 132 deletions(-)
diff --git a/iotdb-core/ainode/ainode/core/config.py
b/iotdb-core/ainode/ainode/core/config.py
index edcdaf559be..678a7cc6d62 100644
--- a/iotdb-core/ainode/ainode/core/config.py
+++ b/iotdb-core/ainode/ainode/core/config.py
@@ -19,6 +19,7 @@ import os
from ainode.core.constant import (
AINODE_BUILD_INFO,
+ AINODE_BUILTIN_MODELS_DIR,
AINODE_CLUSTER_NAME,
AINODE_CONF_DIRECTORY_NAME,
AINODE_CONF_FILE_NAME,
@@ -57,7 +58,7 @@ class AINodeConfig(object):
# Directory to save models
self._ain_models_dir = AINODE_MODELS_DIR
-
+ self._ain_builtin_models_dir = AINODE_BUILTIN_MODELS_DIR
self._ain_system_dir = AINODE_SYSTEM_DIR
# Whether to enable compression for thrift
@@ -130,6 +131,12 @@ class AINodeConfig(object):
def set_ain_models_dir(self, ain_models_dir: str) -> None:
self._ain_models_dir = ain_models_dir
+ def get_ain_builtin_models_dir(self) -> str:
+ return self._ain_builtin_models_dir
+
+ def set_ain_builtin_models_dir(self, ain_builtin_models_dir: str) -> None:
+ self._ain_builtin_models_dir = ain_builtin_models_dir
+
def get_ain_system_dir(self) -> str:
return self._ain_system_dir
@@ -158,7 +165,6 @@ class AINodeConfig(object):
@singleton
class AINodeDescriptor(object):
-
def __init__(self):
self._config = AINodeConfig()
self._load_config_from_file()
diff --git a/iotdb-core/ainode/ainode/core/constant.py
b/iotdb-core/ainode/ainode/core/constant.py
index 98a66abc019..41d0344e056 100644
--- a/iotdb-core/ainode/ainode/core/constant.py
+++ b/iotdb-core/ainode/ainode/core/constant.py
@@ -31,6 +31,7 @@ AINODE_SYSTEM_FILE_NAME = "system.properties"
AINODE_INFERENCE_RPC_ADDRESS = "127.0.0.1"
AINODE_INFERENCE_RPC_PORT = 10810
AINODE_MODELS_DIR = "data/ainode/models"
+AINODE_BUILTIN_MODELS_DIR = "data/ainode/models/weights" # For built-in
models, we only need to store their weights and config.
AINODE_SYSTEM_DIR = "data/ainode/system"
AINODE_LOG_DIR = "logs/ainode"
AINODE_THRIFT_COMPRESSION_ENABLED = False
@@ -53,8 +54,9 @@ AINODE_LOG_FILE_LEVELS = [logging.DEBUG, logging.INFO,
logging.WARNING, logging.
TRIAL_ID_PREFIX = "__trial_"
DEFAULT_TRIAL_ID = TRIAL_ID_PREFIX + "0"
-DEFAULT_MODEL_FILE_NAME = "model.safetensors"
-DEFAULT_CONFIG_FILE_NAME = "config.json"
+
+DEFAULT_MODEL_FILE_NAME = "model.pt"
+DEFAULT_CONFIG_FILE_NAME = "config.yaml"
DEFAULT_CHUNK_SIZE = 8192
DEFAULT_RECONNECT_TIMEOUT = 20
@@ -63,6 +65,12 @@ DEFAULT_RECONNECT_TIMES = 3
STD_LEVEL = logging.INFO
+TIMER_REPO_ID = {
+ "_timerxl": "thuml/timer-base-84m",
+ "_sundial": "thuml/sundial-base-128m",
+}
+
+
class TSStatusCode(Enum):
SUCCESS_STATUS = 200
REDIRECTION_RECOMMEND = 400
@@ -162,6 +170,16 @@ class BuiltInModelType(Enum):
values.append(item.value)
return values
+ @staticmethod
+ def is_built_in_model(model_id: str) -> bool:
+ """
+ Check if the model ID corresponds to a built-in model.
+ """
+ # TODO: Unify this ugly hard code
+ if "timerxl" in model_id or "sundial" in model_id:
+ return True
+ return model_id in BuiltInModelType.values()
+
class AttributeName(Enum):
# forecast Attribute
diff --git a/iotdb-core/ainode/ainode/core/manager/inference_manager.py
b/iotdb-core/ainode/ainode/core/manager/inference_manager.py
index eb8becd0f17..dcbe8bd7359 100644
--- a/iotdb-core/ainode/ainode/core/manager/inference_manager.py
+++ b/iotdb-core/ainode/ainode/core/manager/inference_manager.py
@@ -21,7 +21,7 @@ import pandas as pd
import torch
from iotdb.tsfile.utils.tsblock_serde import deserialize
-from ainode.core.constant import TSStatusCode
+from ainode.core.constant import BuiltInModelType, TSStatusCode
from ainode.core.exception import (
InferenceModelInternalError,
InvalidWindowArgumentError,
@@ -29,6 +29,8 @@ from ainode.core.exception import (
)
from ainode.core.log import Logger
from ainode.core.manager.model_manager import ModelManager
+from ainode.core.model.sundial.modeling_sundial import SundialForPrediction
+from ainode.core.model.timerxl.modeling_timer import TimerForPrediction
from ainode.core.util.serde import convert_to_binary
from ainode.core.util.status import get_status
from ainode.thrift.ainode.ttypes import (
@@ -117,9 +119,9 @@ class RegisteredStrategy(InferenceStrategy):
def _get_strategy(model_id, model):
- if model_id == "_timerxl":
+ if isinstance(model, TimerForPrediction):
return TimerXLStrategy(model)
- if model_id == "_sundial":
+ if isinstance(model, SundialForPrediction):
return SundialStrategy(model)
if model_id.startswith("_"):
return BuiltInStrategy(model)
@@ -127,7 +129,6 @@ def _get_strategy(model_id, model):
class InferenceManager:
-
def __init__(self, model_manager: ModelManager):
self.model_manager = model_manager
@@ -145,18 +146,17 @@ class InferenceManager:
try:
raw = data_getter(req)
full_data = deserializer(raw)
- attrs = extract_attrs(req)
+ inference_attrs = extract_attrs(req)
# load model
- if model_id.startswith("_"):
- model = self.model_manager.load_built_in_model(model_id, attrs)
- else:
- accel = str(attrs.get("acceleration", "")).lower() == "true"
- model = self.model_manager.load_model(model_id, accel)
+ accel = str(inference_attrs.get("acceleration", "")).lower() ==
"true"
+ model = self.model_manager.load_model(
+ model_id, BuiltInModelType.is_built_in_model(model_id), accel
+ )
# inference by strategy
strategy = _get_strategy(model_id, model)
- outputs = strategy.infer(full_data, **attrs)
+ outputs = strategy.infer(full_data, **inference_attrs)
# construct response
status = get_status(TSStatusCode.SUCCESS_STATUS)
diff --git a/iotdb-core/ainode/ainode/core/manager/model_manager.py
b/iotdb-core/ainode/ainode/core/manager/model_manager.py
index ced7277c1ae..b039f924023 100644
--- a/iotdb-core/ainode/ainode/core/manager/model_manager.py
+++ b/iotdb-core/ainode/ainode/core/manager/model_manager.py
@@ -17,16 +17,15 @@
#
from typing import Callable
+from torch import nn
from yaml import YAMLError
-from ainode.core.constant import BuiltInModelType, TSStatusCode
+from ainode.core.constant import TSStatusCode
from ainode.core.exception import (
BadConfigValueError,
- BuiltInModelNotSupportError,
InvalidUriError,
)
from ainode.core.log import Logger
-from ainode.core.model.built_in_model_factory import fetch_built_in_model
from ainode.core.model.model_storage import ModelStorage
from ainode.core.util.status import get_status
from ainode.thrift.ainode.ttypes import (
@@ -96,9 +95,39 @@ class ModelManager:
logger.warning(e)
return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
- def load_model(self, model_id: str, acceleration: bool = False) ->
Callable:
- logger.info(f"load model {model_id}")
- return self.model_storage.load_model(model_id, acceleration)
+ def load_model(
+ self, model_id: str, is_built_in: bool, acceleration: bool = False
+ ) -> Callable:
+ """
+ Load the model with the given model_id.
+ """
+ logger.info(f"Load model {model_id}")
+ try:
+ model = self.model_storage.load_model(
+ model_id.lower(), is_built_in, acceleration
+ )
+ logger.info(f"Model {model_id} loaded")
+ return model
+ except Exception as e:
+ logger.error(f"Failed to load model {model_id}: {e}")
+ raise
+
+ def save_model(
+ self, model_id: str, is_built_in: bool, model: nn.Module
+ ) -> TSStatus:
+ """
+ Save the model using save_pretrained
+ """
+ logger.info(f"Saving model {model_id}")
+ try:
+ self.model_storage.save_model(model_id, is_built_in, model)
+ logger.info(f"Saving model {model_id} successfully")
+ return get_status(
+ TSStatusCode.SUCCESS_STATUS, f"Model {model_id} saved
successfully"
+ )
+ except Exception as e:
+ logger.error(f"Save model failed: {e}")
+ return get_status(TSStatusCode.AINODE_INTERNAL_ERROR, str(e))
def get_ckpt_path(self, model_id: str) -> str:
"""
@@ -111,10 +140,3 @@ class ModelManager:
str: The path to the checkpoint file for the model.
"""
return self.model_storage.get_ckpt_path(model_id)
-
- @staticmethod
- def load_built_in_model(model_id: str, attributes: {}):
- model_id = model_id.lower()
- if model_id not in BuiltInModelType.values():
- raise BuiltInModelNotSupportError(model_id)
- return fetch_built_in_model(model_id, attributes)
diff --git a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py
b/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py
index 8bd3bfc4800..b6e2de7115d 100644
--- a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py
+++ b/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py
@@ -17,9 +17,10 @@
#
import os
from abc import abstractmethod
-from typing import Dict, List
+from typing import Callable, Dict, List
import numpy as np
+from huggingface_hub import hf_hub_download
from sklearn.preprocessing import MinMaxScaler
from sktime.annotation.hmm_learn import GMMHMM, GaussianHMM
from sktime.annotation.stray import STRAY
@@ -29,7 +30,7 @@ from sktime.forecasting.naive import NaiveForecaster
from sktime.forecasting.trend import STLForecaster
from ainode.core.config import AINodeDescriptor
-from ainode.core.constant import AttributeName, BuiltInModelType
+from ainode.core.constant import TIMER_REPO_ID, AttributeName, BuiltInModelType
from ainode.core.exception import (
BuiltInModelNotSupportError,
InferenceModelInternalError,
@@ -40,13 +41,43 @@ from ainode.core.exception import (
)
from ainode.core.log import Logger
from ainode.core.model.sundial import modeling_sundial
-from ainode.core.model.sundial.configuration_sundial import SundialConfig
from ainode.core.model.timerxl import modeling_timer
-from ainode.core.model.timerxl.configuration_timer import TimerConfig
logger = Logger()
+def download_built_in_model_if_necessary(model_id: str, local_dir):
+ """
+ Download the built-in model from HuggingFace repository when necessary.
+ """
+ if "_timer" == model_id or "_sundial" == model_id:
+ weights_path = os.path.join(local_dir, "model.safetensors")
+ if not os.path.exists(weights_path):
+ logger.info(
+ f"Weight not found at {weights_path}, downloading from
HuggingFace..."
+ )
+ repo_id = TIMER_REPO_ID[model_id]
+ try:
+ hf_hub_download(
+ repo_id=repo_id,
+ filename="model.safetensors",
+ local_dir=local_dir,
+ )
+ logger.info(f"Got weight to {weights_path}")
+ config_path = os.path.join(local_dir, "config.json")
+ hf_hub_download(
+ repo_id=repo_id,
+ filename="config.json",
+ local_dir=local_dir,
+ )
+ logger.info(f"Got config to {config_path}")
+ except Exception as e:
+ logger.error(
+ f"Failed to download huggingface model to {local_dir} due
to {e}"
+ )
+ raise e
+
+
def get_model_attributes(model_id: str):
if model_id == BuiltInModelType.ARIMA.value:
attribute_map = arima_attribute_map
@@ -65,34 +96,26 @@ def get_model_attributes(model_id: str):
attribute_map = gaussian_hmm_attribute_map
elif model_id == BuiltInModelType.STRAY.value:
attribute_map = stray_attribute_map
- elif model_id == BuiltInModelType.TIMER_XL.value:
+ # TODO: The model type should be judged before enter this file
+ elif "timerxl" in model_id:
attribute_map = timerxl_attribute_map
- elif model_id == BuiltInModelType.SUNDIAL.value:
+ elif "sundial" in model_id:
attribute_map = sundial_attribute_map
else:
raise BuiltInModelNotSupportError(model_id)
return attribute_map
-def fetch_built_in_model(model_id, inference_attributes):
+def fetch_built_in_model(model_id: str, model_dir) -> Callable:
"""
+ Fetch the built-in model according to its id and directory, not that this
directory only contains model weights and config.
Args:
model_id: the unique id of the model
- inference_attributes: a list of attributes to be inferred, in this
function, the attributes will include some
- parameters of the built-in model. Some parameters are optional,
and if the parameters are not
- specified, the default value will be used.
+ model_dir: for huggingface models only, the directory where the model
is stored
Returns:
model: the built-in model
- attributes: a dict of attributes, where the key is the attribute name,
the value is the parsed value of the
- attribute
- Description:
- the create_built_in_model function will create the built-in model,
which does not require user
- registration. This module will parse the inference attributes and
create the built-in model.
"""
- attribute_map = get_model_attributes(model_id)
-
- # parse the inference attributes, attributes is a Dict[str, Any]
- attributes = parse_attribute(inference_attributes, attribute_map)
+ attributes = get_model_attributes(model_id)
# build the built-in model
if model_id == BuiltInModelType.ARIMA.value:
@@ -113,11 +136,9 @@ def fetch_built_in_model(model_id, inference_attributes):
elif model_id == BuiltInModelType.STRAY.value:
model = STRAYModel(attributes)
elif model_id == BuiltInModelType.TIMER_XL.value:
- model =
modeling_timer.TimerForPrediction(TimerConfig.from_dict(attributes))
+ model = modeling_timer.TimerForPrediction.from_pretrained(model_dir)
elif model_id == BuiltInModelType.SUNDIAL.value:
- model = modeling_sundial.SundialForPrediction(
- SundialConfig.from_dict(attributes)
- )
+ model =
modeling_sundial.SundialForPrediction.from_pretrained(model_dir)
else:
raise BuiltInModelNotSupportError(model_id)
diff --git a/iotdb-core/ainode/ainode/core/model/model_storage.py
b/iotdb-core/ainode/ainode/core/model/model_storage.py
index 864b5c30e0a..a5535df56c4 100644
--- a/iotdb-core/ainode/ainode/core/model/model_storage.py
+++ b/iotdb-core/ainode/ainode/core/model/model_storage.py
@@ -20,14 +20,23 @@ import os
import shutil
from collections.abc import Callable
-import torch
-import torch._dynamo
from pylru import lrucache
+from torch import nn
from ainode.core.config import AINodeDescriptor
-from ainode.core.constant import DEFAULT_CONFIG_FILE_NAME,
DEFAULT_MODEL_FILE_NAME
-from ainode.core.exception import ModelNotExistError
+from ainode.core.constant import (
+ DEFAULT_CONFIG_FILE_NAME,
+ DEFAULT_MODEL_FILE_NAME,
+ BuiltInModelType,
+)
+from ainode.core.exception import (
+ BuiltInModelNotSupportError,
+)
from ainode.core.log import Logger
+from ainode.core.model.built_in_model_factory import (
+ download_built_in_model_if_necessary,
+ fetch_built_in_model,
+)
from ainode.core.model.model_factory import fetch_model_by_uri
from ainode.core.util.lock import ModelLockPool
@@ -45,6 +54,15 @@ class ModelStorage(object):
except PermissionError as e:
logger.error(e)
raise e
+ self._builtin_model_dir = os.path.join(
+ os.getcwd(),
AINodeDescriptor().get_config().get_ain_builtin_models_dir()
+ )
+ if not os.path.exists(self._builtin_model_dir):
+ try:
+ os.makedirs(self._builtin_model_dir)
+ except PermissionError as e:
+ logger.error(e)
+ raise e
self._lock_pool = ModelLockPool()
self._model_cache = lrucache(
AINodeDescriptor().get_config().get_ain_model_storage_cache_size()
@@ -68,40 +86,6 @@ class ModelStorage(object):
config_storage_path = os.path.join(storage_path,
DEFAULT_CONFIG_FILE_NAME)
return fetch_model_by_uri(uri, model_storage_path, config_storage_path)
- def load_model(self, model_id: str, acceleration: bool) -> Callable:
- """
- Returns:
- model: a ScriptModule contains model architecture and parameters,
which can be deployed cross-platform
- """
- ain_models_dir = os.path.join(self._model_dir, f"{model_id}")
- model_path = os.path.join(ain_models_dir, DEFAULT_MODEL_FILE_NAME)
- with self._lock_pool.get_lock(model_id).read_lock():
- if model_path in self._model_cache:
- model = self._model_cache[model_path]
- if (
- isinstance(model, torch._dynamo.eval_frame.OptimizedModule)
- or not acceleration
- ):
- return model
- else:
- model = torch.compile(model)
- self._model_cache[model_path] = model
- return model
- else:
- if not os.path.exists(model_path):
- raise ModelNotExistError(model_path)
- else:
- model = torch.jit.load(model_path)
- if acceleration:
- try:
- model = torch.compile(model)
- except Exception as e:
- logger.warning(
- f"acceleration failed, fallback to normal
mode: {str(e)}"
- )
- self._model_cache[model_path] = model
- return model
-
def delete_model(self, model_id: str) -> None:
"""
Args:
@@ -120,6 +104,46 @@ class ModelStorage(object):
if file_path in self._model_cache:
del self._model_cache[file_path]
+ def load_model(
+ self, model_id: str, is_built_in: bool, acceleration: bool
+ ) -> Callable:
+ """
+ Load a model with automatic detection of .safetensors or .pt format
+
+ Returns:
+ model: The model instance corresponding to specific model_id
+ """
+ with self._lock_pool.get_lock(model_id).read_lock():
+ if is_built_in:
+ if model_id not in BuiltInModelType.values():
+ raise BuiltInModelNotSupportError(model_id)
+ # For built-in models, we support auto download
+ model_dir = os.path.join(self._builtin_model_dir,
f"{model_id}")
+ download_built_in_model_if_necessary(model_id, model_dir)
+ return fetch_built_in_model(model_id, model_dir)
+ else:
+ # TODO: support load the user-defined model
+ # model_dir = os.path.join(self._model_dir, f"{model_id}")
+ raise NotImplementedError
+
+ def save_model(self, model_id: str, is_built_in: bool, model: nn.Module):
+ """
+ Save the model using save_pretrained
+
+ Returns:
+ Whether saving succeeded
+ """
+ with self._lock_pool.get_lock(model_id).write_lock():
+ if is_built_in:
+ if model_id not in BuiltInModelType.values():
+ raise BuiltInModelNotSupportError(model_id)
+ model_dir = os.path.join(self._builtin_model_dir,
f"{model_id}")
+ model.save_pretrained(model_dir)
+ else:
+ # TODO: support save the user-defined model
+ # model_dir = os.path.join(self._model_dir, f"{model_id}")
+ raise NotImplementedError
+
def get_ckpt_path(self, model_id: str) -> str:
"""
Get the checkpoint path for a given model ID.
@@ -130,4 +154,5 @@ class ModelStorage(object):
Returns:
str: The path to the checkpoint file for the model.
"""
- return os.path.join(self._model_dir, f"{model_id}")
+ # Only support built-in models for now
+ return os.path.join(self._builtin_model_dir, f"{model_id}")
diff --git a/iotdb-core/ainode/ainode/core/model/sundial/modeling_sundial.py
b/iotdb-core/ainode/ainode/core/model/sundial/modeling_sundial.py
index a74e8e6cf23..2e6b436fbda 100644
--- a/iotdb-core/ainode/ainode/core/model/sundial/modeling_sundial.py
+++ b/iotdb-core/ainode/ainode/core/model/sundial/modeling_sundial.py
@@ -471,27 +471,7 @@ class SundialForPrediction(SundialPreTrainedModel,
TSGenerationMixin):
self.config.hidden_size,
self.config.num_sampling_steps,
)
- # TODO: Unify data loader
- if not os.path.exists(config.ckpt_path):
- os.mkdir(config.ckpt_path)
- weights_path = os.path.join(config.ckpt_path, "model.safetensors")
- if not os.path.exists(weights_path):
- logger.info(
- f"Weight not found at {weights_path}, downloading from
HuggingFace..."
- )
- repo_id = "thuml/sundial-base-128m"
- try:
- hf_hub_download(
- repo_id=repo_id,
- filename="model.safetensors",
- local_dir=config.ckpt_path,
- )
- logger.info(f"Got weight to {weights_path}")
- except Exception as e:
- logger.error(f"Failed to download weight to {weights_path} due
to {e}")
- raise e
- state_dict = load_safetensors(weights_path)
- self.load_state_dict(state_dict, strict=True)
+ self.post_init()
def set_decoder(self, decoder):
self.model = decoder
diff --git a/iotdb-core/ainode/ainode/core/model/timerxl/modeling_timer.py
b/iotdb-core/ainode/ainode/core/model/timerxl/modeling_timer.py
index 42b3a82b972..42566c0e1c9 100644
--- a/iotdb-core/ainode/ainode/core/model/timerxl/modeling_timer.py
+++ b/iotdb-core/ainode/ainode/core/model/timerxl/modeling_timer.py
@@ -467,27 +467,7 @@ class TimerForPrediction(TimerPreTrainedModel,
TSGenerationMixin):
self.output_token_len_map[output_token_len] = i
self.lm_heads = nn.ModuleList(lm_head_list)
self.loss_function = torch.nn.MSELoss(reduction="none")
- # TODO: Unify data loader
- if not os.path.exists(config.ckpt_path):
- os.mkdir(config.ckpt_path)
- weights_path = os.path.join(config.ckpt_path, "model.safetensors")
- if not os.path.exists(weights_path):
- logger.info(
- f"Weight not found at {weights_path}, downloading from
HuggingFace..."
- )
- repo_id = "thuml/sundial-base-128m"
- try:
- hf_hub_download(
- repo_id=repo_id,
- filename="model.safetensors",
- local_dir=config.ckpt_path,
- )
- logger.info(f"Got weight to {weights_path}")
- except Exception as e:
- logger.error(f"Failed to download weight to {weights_path} due
to {e}")
- raise e
- state_dict = load_safetensors(weights_path)
- self.load_state_dict(state_dict, strict=True)
+ self.post_init()
def set_decoder(self, decoder):
self.model = decoder
