Reconstruct model management and model loading

Author: RkGrit

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

@@ -30,21 +30,21 @@
     AINODE_CONF_FILE_NAME,
     AINODE_CONF_GIT_FILE_NAME,
     AINODE_CONF_POM_FILE_NAME,
+    AINODE_FINETUNE_MODELS_DIR,
     AINODE_INFERENCE_BATCH_INTERVAL_IN_MS,
     AINODE_INFERENCE_EXTRA_MEMORY_RATIO,
     AINODE_INFERENCE_MAX_PREDICT_LENGTH,
     AINODE_INFERENCE_MEMORY_USAGE_RATIO,
     AINODE_INFERENCE_MODEL_MEM_USAGE_MAP,
     AINODE_LOG_DIR,
     AINODE_MODELS_DIR,
-    AINODE_ROOT_CONF_DIRECTORY_NAME,
-    AINODE_ROOT_DIR,
     AINODE_RPC_ADDRESS,
     AINODE_RPC_PORT,
     AINODE_SYSTEM_DIR,
     AINODE_SYSTEM_FILE_NAME,
     AINODE_TARGET_CONFIG_NODE_LIST,
     AINODE_THRIFT_COMPRESSION_ENABLED,
+    AINODE_USER_DEFINED_MODELS_DIR,
     AINODE_VERSION_INFO,
 )
 from iotdb.ainode.core.exception import BadNodeUrlError
@@ -97,6 +97,8 @@ def __init__(self):
         # Directory to save models
         self._ain_models_dir = AINODE_MODELS_DIR
         self._ain_builtin_models_dir = AINODE_BUILTIN_MODELS_DIR
+        self._ain_finetune_models_dir = AINODE_FINETUNE_MODELS_DIR
+        self._ain_user_defined_models_dir = AINODE_USER_DEFINED_MODELS_DIR
         self._ain_system_dir = AINODE_SYSTEM_DIR
 
         # Whether to enable compression for thrift
@@ -211,6 +213,18 @@ def get_ain_builtin_models_dir(self) -> str:
     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_finetune_models_dir(self) -> str:
+        return self._ain_finetune_models_dir
+
+    def set_ain_finetune_models_dir(self, ain_finetune_models_dir: str) -> None:
+        self._ain_finetune_models_dir = ain_finetune_models_dir
+
+    def get_ain_user_defined_models_dir(self) -> str:
+        return self._ain_user_defined_models_dir
+
+    def set_ain_user_defined_models_dir(self, ain_user_defined_models_dir: str) -> None:
+        self._ain_user_defined_models_dir = ain_user_defined_models_dir
+
     def get_ain_system_dir(self) -> str:
         return self._ain_system_dir
 
@@ -315,9 +329,7 @@ def _load_config_from_file(self) -> None:
             if "ainode_id" in system_configs:
                 self._config.set_ainode_id(int(system_configs["ainode_id"]))
 
-        git_file = os.path.join(
-            AINODE_ROOT_DIR, AINODE_ROOT_CONF_DIRECTORY_NAME, AINODE_CONF_GIT_FILE_NAME
-        )
+        git_file = os.path.join(AINODE_CONF_DIRECTORY_NAME, AINODE_CONF_GIT_FILE_NAME)
         if os.path.exists(git_file):
             git_configs = load_properties(git_file)
             if "git.commit.id.abbrev" in git_configs:
@@ -327,9 +339,7 @@ def _load_config_from_file(self) -> None:
                         build_info += "-dev"
                 self._config.set_build_info(build_info)
 
-        pom_file = os.path.join(
-            AINODE_ROOT_DIR, AINODE_ROOT_CONF_DIRECTORY_NAME, AINODE_CONF_POM_FILE_NAME
-        )
+        pom_file = os.path.join(AINODE_CONF_DIRECTORY_NAME, AINODE_CONF_POM_FILE_NAME)
         if os.path.exists(pom_file):
             pom_configs = load_properties(pom_file)
             if "version" in pom_configs:

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

@@ -21,13 +21,12 @@
 from enum import Enum
 from typing import List
 
-from iotdb.ainode.core.model.model_enums import BuiltInModelType
 from iotdb.thrift.common.ttypes import TEndPoint
 
+IOTDB_AINODE_HOME = os.getenv("IOTDB_AINODE_HOME", "")
 AINODE_VERSION_INFO = "UNKNOWN"
 AINODE_BUILD_INFO = "UNKNOWN"
-AINODE_CONF_DIRECTORY_NAME = "conf"
-AINODE_ROOT_CONF_DIRECTORY_NAME = "conf"
+AINODE_CONF_DIRECTORY_NAME = os.path.join(IOTDB_AINODE_HOME, "conf")
 AINODE_CONF_FILE_NAME = "iotdb-ainode.properties"
 AINODE_CONF_GIT_FILE_NAME = "git.properties"
 AINODE_CONF_POM_FILE_NAME = "pom.properties"
@@ -53,22 +52,27 @@
 AINODE_INFERENCE_BATCH_INTERVAL_IN_MS = 15
 AINODE_INFERENCE_MAX_PREDICT_LENGTH = 2880
 AINODE_INFERENCE_MODEL_MEM_USAGE_MAP = {
-    BuiltInModelType.SUNDIAL.value: 1036 * 1024**2,  # 1036 MiB
-    BuiltInModelType.TIMER_XL.value: 856 * 1024**2,  # 856 MiB
+    "sundial": 1036 * 1024**2,  # 1036 MiB
+    "timerxl": 856 * 1024**2,  # 856 MiB
 }  # the memory usage of each model in bytes
 AINODE_INFERENCE_MEMORY_USAGE_RATIO = 0.4  # the device space allocated for inference
 AINODE_INFERENCE_EXTRA_MEMORY_RATIO = (
     1.2  # the overhead ratio for inference, used to estimate the pool size
 )
 
-# AINode folder structure
-AINODE_ROOT_DIR = os.path.dirname(
-    os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
+AINODE_MODELS_DIR = os.path.join(IOTDB_AINODE_HOME, "data/ainode/models")
+AINODE_BUILTIN_MODELS_DIR = os.path.join(
+    IOTDB_AINODE_HOME, "data/ainode/models/builtin"
+)  # For built-in models, we only need to store their weights and config.
+AINODE_FINETUNE_MODELS_DIR = os.path.join(
+    IOTDB_AINODE_HOME, "data/ainode/models/finetune"
+)
+AINODE_USER_DEFINED_MODELS_DIR = os.path.join(
+    IOTDB_AINODE_HOME, "data/ainode/models/user_defined"
 )
-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_CACHE_DIR = os.path.expanduser("~/.cache/ainode")
 
 # AINode log
 LOG_FILE_TYPE = ["all", "info", "warn", "error"]
@@ -143,132 +147,8 @@ def name(self):
         return self.value
 
 
-class ForecastModelType(Enum):
-    DLINEAR = "dlinear"
-    DLINEAR_INDIVIDUAL = "dlinear_individual"
-    NBEATS = "nbeats"
-
-    @classmethod
-    def values(cls) -> List[str]:
-        values = []
-        for item in list(cls):
-            values.append(item.value)
-        return values
-
-
 class ModelInputName(Enum):
     DATA_X = "data_x"
     TIME_STAMP_X = "time_stamp_x"
     TIME_STAMP_Y = "time_stamp_y"
     DEC_INP = "dec_inp"
-
-
-class AttributeName(Enum):
-    # forecast Attribute
-    PREDICT_LENGTH = "predict_length"
-
-    # NaiveForecaster
-    STRATEGY = "strategy"
-    SP = "sp"
-
-    # STLForecaster
-    # SP = 'sp'
-    SEASONAL = "seasonal"
-    SEASONAL_DEG = "seasonal_deg"
-    TREND_DEG = "trend_deg"
-    LOW_PASS_DEG = "low_pass_deg"
-    SEASONAL_JUMP = "seasonal_jump"
-    TREND_JUMP = "trend_jump"
-    LOSS_PASS_JUMP = "low_pass_jump"
-
-    # ExponentialSmoothing
-    DAMPED_TREND = "damped_trend"
-    INITIALIZATION_METHOD = "initialization_method"
-    OPTIMIZED = "optimized"
-    REMOVE_BIAS = "remove_bias"
-    USE_BRUTE = "use_brute"
-
-    # Arima
-    ORDER = "order"
-    SEASONAL_ORDER = "seasonal_order"
-    METHOD = "method"
-    MAXITER = "maxiter"
-    SUPPRESS_WARNINGS = "suppress_warnings"
-    OUT_OF_SAMPLE_SIZE = "out_of_sample_size"
-    SCORING = "scoring"
-    WITH_INTERCEPT = "with_intercept"
-    TIME_VARYING_REGRESSION = "time_varying_regression"
-    ENFORCE_STATIONARITY = "enforce_stationarity"
-    ENFORCE_INVERTIBILITY = "enforce_invertibility"
-    SIMPLE_DIFFERENCING = "simple_differencing"
-    MEASUREMENT_ERROR = "measurement_error"
-    MLE_REGRESSION = "mle_regression"
-    HAMILTON_REPRESENTATION = "hamilton_representation"
-    CONCENTRATE_SCALE = "concentrate_scale"
-
-    # GAUSSIAN_HMM
-    N_COMPONENTS = "n_components"
-    COVARIANCE_TYPE = "covariance_type"
-    MIN_COVAR = "min_covar"
-    STARTPROB_PRIOR = "startprob_prior"
-    TRANSMAT_PRIOR = "transmat_prior"
-    MEANS_PRIOR = "means_prior"
-    MEANS_WEIGHT = "means_weight"
-    COVARS_PRIOR = "covars_prior"
-    COVARS_WEIGHT = "covars_weight"
-    ALGORITHM = "algorithm"
-    N_ITER = "n_iter"
-    TOL = "tol"
-    PARAMS = "params"
-    INIT_PARAMS = "init_params"
-    IMPLEMENTATION = "implementation"
-
-    # GMMHMM
-    # N_COMPONENTS = "n_components"
-    N_MIX = "n_mix"
-    # MIN_COVAR = "min_covar"
-    # STARTPROB_PRIOR = "startprob_prior"
-    # TRANSMAT_PRIOR = "transmat_prior"
-    WEIGHTS_PRIOR = "weights_prior"
-
-    # MEANS_PRIOR = "means_prior"
-    # MEANS_WEIGHT = "means_weight"
-    # ALGORITHM = "algorithm"
-    # COVARIANCE_TYPE = "covariance_type"
-    # N_ITER = "n_iter"
-    # TOL = "tol"
-    # INIT_PARAMS = "init_params"
-    # PARAMS = "params"
-    # IMPLEMENTATION = "implementation"
-
-    # STRAY
-    ALPHA = "alpha"
-    K = "k"
-    KNN_ALGORITHM = "knn_algorithm"
-    P = "p"
-    SIZE_THRESHOLD = "size_threshold"
-    OUTLIER_TAIL = "outlier_tail"
-
-    # timerxl
-    INPUT_TOKEN_LEN = "input_token_len"
-    HIDDEN_SIZE = "hidden_size"
-    INTERMEDIATE_SIZE = "intermediate_size"
-    OUTPUT_TOKEN_LENS = "output_token_lens"
-    NUM_HIDDEN_LAYERS = "num_hidden_layers"
-    NUM_ATTENTION_HEADS = "num_attention_heads"
-    HIDDEN_ACT = "hidden_act"
-    USE_CACHE = "use_cache"
-    ROPE_THETA = "rope_theta"
-    ATTENTION_DROPOUT = "attention_dropout"
-    INITIALIZER_RANGE = "initializer_range"
-    MAX_POSITION_EMBEDDINGS = "max_position_embeddings"
-    CKPT_PATH = "ckpt_path"
-
-    # sundial
-    DROPOUT_RATE = "dropout_rate"
-    FLOW_LOSS_DEPTH = "flow_loss_depth"
-    NUM_SAMPLING_STEPS = "num_sampling_steps"
-    DIFFUSION_BATCH_MUL = "diffusion_batch_mul"
-
-    def name(self) -> str:
-        return self.value

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

@@ -15,14 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+
 import threading
 from typing import Any
 
 import torch
 
-from iotdb.ainode.core.inference.strategy.abstract_inference_pipeline import (
-    AbstractInferencePipeline,
-)
 from iotdb.ainode.core.log import Logger
 from iotdb.ainode.core.util.atmoic_int import AtomicInt
 
@@ -41,7 +39,6 @@ def __init__(
         req_id: str,
         model_id: str,
         inputs: torch.Tensor,
-        inference_pipeline: AbstractInferencePipeline,
         max_new_tokens: int = 96,
         **infer_kwargs,
     ):
@@ -52,7 +49,6 @@ def __init__(
         self.model_id = model_id
         self.inputs = inputs
         self.infer_kwargs = infer_kwargs
-        self.inference_pipeline = inference_pipeline
         self.max_new_tokens = (
             max_new_tokens  # Number of time series data points to generate
         )