Refactoring of Model Storage, Loading, and Inference Pipeline #16794

RkGrit · 2025-11-24T14:25:00Z

This PR introduces significant improvements in the model storage, loading, and inference pipeline management for better extensibility, efficiency, and ease of use. The changes include the refactoring of model storage to support a wider range of models, streamlining the model loading process, and the introduction of a unified inference pipeline. These improvements aim to optimize model management, reduce memory usage, and enhance the overall inference workflow.

Model Storage Refactoring
- Extended Support for Models: The system now supports not only built-in models like TimerXL and Sundial but also allows the integration of fine-tuned and user-defined models.
- Unified Model Management: A new model management system enables model registration, deletion, and loading from both local paths and Hugging Face.
- Code Optimization: Redundant code from previous versions has been removed, and hard-coded model management has been replaced by a more flexible approach that integrates seamlessly with the Hugging Face Transformers ecosystem.
Model Loading Refactoring
- Simplified Model Loading: The previous custom loading logic with complex if...else... conditions has been replaced by a unified model loading interface, simplifying the process.
- Automatic Model Type Detection: The system now automatically detects the model type and selects the appropriate loading method, supporting models from Transformers, sktime, and PyTorch.
- Lazy Loading: The PR introduces lazy loading for Python modules, eliminating the need to load multiple modules at startup, reducing initialization time and memory consumption.
Inference Pipeline Addition
- Unified Inference Workflow: The introduction of the Inference Pipeline encapsulates the entire model inference process, offering a standardized interface for preprocessing, inference, and post-processing.
- Support for Multiple Tasks: The pipeline is versatile, supporting various inference tasks such as prediction, classification, and dialogue-based tasks.

Copilot

Pull request overview

This PR significantly refactors the model storage, loading, and inference pipeline architecture to improve extensibility and maintainability. The changes transition from hard-coded model management to a flexible, discovery-based system that seamlessly integrates with HuggingFace Transformers and sktime.

Key Changes

Unified Model Storage: Introduced a category-based storage system (builtin/user_defined/finetune) with automatic model discovery and lazy registration for Transformers models
Simplified Model Loading: Replaced complex conditional logic with a unified ModelLoader class that automatically detects model types (Transformers, sktime, PyTorch) and applies appropriate loading strategies
Inference Pipeline Framework: Created a modular pipeline architecture with base classes (BasicPipeline, ForecastPipeline) and model-specific implementations for timerxl and sundial models

Reviewed changes

Copilot reviewed 35 out of 37 changed files in this pull request and generated 12 comments.

Show a summary per file

File	Description
model_storage.py	Complete rewrite to support category-based storage with discovery, registration, and lazy Transformers model registration
model_loader.py	New unified loader supporting Transformers, sktime, and PyTorch models with automatic type detection
model_info.py	Simplified model info with updated enums, removed complex type detection, added REPO_ID_MAP for HF downloads
model_enums.py	Refactored enums: removed BuiltInModelType, updated ModelCategory values, added UriType
handler.py	Updated to use `get_model_manager()` singleton, inlined validation logic, improved error messages
model_manager.py	Refactored to use composition (storage + loader), added discovery on init, new query methods
inference_manager.py	Simplified strategy pattern, removed model-specific inference classes, uses new pipeline system
pipeline/*.py	New modular pipeline architecture with basic, forecast, classification, and chat pipelines
sktime/*.py	New sktime model support with configuration system and model factory
poetry.lock	Added platform-specific greenlet and statsmodels wheels for Python 3.14

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review for a chance to win a $100 gift card. Take the survey.

Copilot · 2025-11-25T03:08:11Z

iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py

+            "covariance_type",
+            "diag",
+            "str",
+            choices=["sperical", "diag", "full", "tied"],


Typo in choice value "sperical" should be "spherical". This will cause validation failures when users try to use the correct spelling.

Suggested change

choices=["sperical", "diag", "full", "tied"],

choices=["spherical", "diag", "full", "tied"],

Copilot · 2025-11-25T03:08:11Z

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

 def estimate_pool_size(device: torch.device, model_id: str) -> int:
-    model_info = BUILT_IN_LTSM_MAP.get(model_id, None)
-    if model_info is None or model_info.model_type not in MODEL_MEM_USAGE_MAP:
+    model_info = get_model_manager.get_model_info(model_id)


Missing closing parenthesis. The function call get_model_manager.get_model_info(model_id) is missing parentheses after get_model_manager. It should be get_model_manager().get_model_info(model_id).

Suggested change

model_info = get_model_manager.get_model_info(model_id)

model_info = get_model_manager().get_model_info(model_id)

Copilot · 2025-11-25T03:08:11Z

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

+    def _process_model_directory(
+        self, model_dir: Path, model_id: str, category: ModelCategory
+    ):
+        """Handling the discovery logic for a single model directory."""


[nitpick] Missing period at end of comment. The comment should end with a period for consistency with documentation standards.

Copilot · 2025-11-25T03:08:12Z

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

 import concurrent.futures
-import json
 import os
 import shutil
-from collections.abc import Callable
-from typing import Dict
+from typing import List, Optional

-import torch
-from torch import nn
+from huggingface_hub import hf_hub_download, snapshot_download
+from transformers import AutoConfig, AutoModelForCausalLM

-from iotdb.ainode.core.config import AINodeDescriptor
-from iotdb.ainode.core.constant import (
-    MODEL_CONFIG_FILE_IN_JSON,
-    MODEL_WEIGHTS_FILE_IN_PT,
-    TSStatusCode,
-)
-from iotdb.ainode.core.exception import (
-    BuiltInModelDeletionError,
-    ModelNotExistError,
-    UnsupportedError,
-)
+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.built_in_model_factory import (
-    download_built_in_ltsm_from_hf_if_necessary,
-    fetch_built_in_model,
-)
-from iotdb.ainode.core.model.model_enums import (
-    BuiltInModelType,
-    ModelCategory,
-    ModelFileType,
-    ModelStates,
-)
-from iotdb.ainode.core.model.model_factory import fetch_model_by_uri
+from iotdb.ainode.core.model.model_enums import REPO_ID_MAP, ModelCategory, ModelStates
 from iotdb.ainode.core.model.model_info import (
-    BUILT_IN_LTSM_MAP,
-    BUILT_IN_MACHINE_LEARNING_MODEL_MAP,
+    BUILTIN_HF_TRANSFORMERS_MODEL_MAP,
+    BUILTIN_SKTIME_MODEL_MAP,
    ModelInfo,
-    get_built_in_model_type,
-    get_model_file_type,
 )
-from iotdb.ainode.core.model.uri_utils import get_model_register_strategy
+from iotdb.ainode.core.model.utils import *


Missing import statement. The code uses json module on line 206 but doesn't import it at the top of the file. Add import json to the imports section.

Copilot · 2025-11-25T03:08:12Z

iotdb-core/ainode/iotdb/ainode/core/model/sktime/configuration_sktime.py

+        "pipeline": AttributeConfig(
+            "pipeline", "last", "str", choices=["last", "mean"]


The attribute name "pipeline" should likely be "strategy" to match the NaiveForecaster parameter. According to sktime documentation, the NaiveForecaster parameter for choosing between "last" and "mean" is called "strategy", not "pipeline".

Suggested change

"pipeline": AttributeConfig(

"pipeline", "last", "str", choices=["last", "mean"]

"strategy": AttributeConfig(

"strategy", "last", "str", choices=["last", "mean"]

Copilot · 2025-11-25T03:08:13Z

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

-    ModelFileType,
-    ModelStates,
-)
+from typing import Dict, List, Optional, Tuple


Import of 'List' is not used.
Import of 'Tuple' is not used.

Suggested change

from typing import Dict, List, Optional, Tuple

from typing import Dict, Optional

Copilot · 2025-11-25T03:08:13Z

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

+# under the License.
+#
+
+import pandas as pd


Import of 'pd' is not used.

Suggested change

import pandas as pd

Copilot · 2025-11-25T03:08:14Z

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

+import torch
+
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
+from iotdb.ainode.core.util.serde import convert_to_binary


Import of 'convert_to_binary' is not used.

Suggested change

from iotdb.ainode.core.util.serde import convert_to_binary

Copilot · 2025-11-25T03:08:14Z

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

+# under the License.
+#
+
+import pandas as pd


Import of 'pd' is not used.

Suggested change

import pandas as pd

Copilot · 2025-11-25T03:08:14Z

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

+import torch
+
+from iotdb.ainode.core.inference.pipeline.basic_pipeline import ForecastPipeline
+from iotdb.ainode.core.util.serde import convert_to_binary


Import of 'convert_to_binary' is not used.

Suggested change

from iotdb.ainode.core.util.serde import convert_to_binary

SpriCoder

Great Jobs. Please take a look of CI

CRZbulabula

PTAL.

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

CRZbulabula · 2025-11-25T11:27:20Z

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

+    if model_id == "timerxl":
+        return TimerxlPipeline(model_id, device=device)
+    elif model_id == "sundial":
+        return SundialPipeline(model_id, device=device)


Under current implementation, to integrate a new model, we should add another if-else branch here? Is it toooooo much?

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

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

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

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

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

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

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

CRZbulabula

PTAL.

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

CRZbulabula · 2025-11-26T05:46:37Z

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

+            batch_output = self._inference_pipeline.infer(
+                batch_inputs,
+                predict_length=requests[0].max_new_tokens,
+                # num_samples=10,
+                revin=True,
+            )


In your current implementation, there is no parameters can be delivered into inference_pipeline, we should further discuss.

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

RkGrit added 2 commits November 10, 2025 17:20

refactor_built_in_models

0c22906

delete old code in model folder

de9fe48

RkGrit force-pushed the model_management_v2 branch from ff6fa25 to 2ade43f Compare November 24, 2025 14:34

Reconstruct model management and model loading

18cccc3

RkGrit force-pushed the model_management_v2 branch from 2ade43f to 18cccc3 Compare November 24, 2025 14:37

Merge branch 'master' into model_management_v2

92d4643

CRZbulabula requested a review from Copilot November 25, 2025 03:03

Copilot started reviewing on behalf of CRZbulabula November 25, 2025 03:05 View session

Copilot finished reviewing on behalf of CRZbulabula November 25, 2025 03:06

Copilot AI reviewed Nov 25, 2025

View reviewed changes

SpriCoder requested changes Nov 25, 2025

View reviewed changes

CRZbulabula requested changes Nov 25, 2025

View reviewed changes

CRZbulabula requested changes Nov 26, 2025

View reviewed changes

	choices=["sperical", "diag", "full", "tied"],
	choices=["spherical", "diag", "full", "tied"],

	model_info = get_model_manager.get_model_info(model_id)
	model_info = get_model_manager().get_model_info(model_id)

		"pipeline": AttributeConfig(
		"pipeline", "last", "str", choices=["last", "mean"]

	from typing import Dict, List, Optional, Tuple
	from typing import Dict, Optional

Refactoring of Model Storage, Loading, and Inference Pipeline #16794

Are you sure you want to change the base?

Refactoring of Model Storage, Loading, and Inference Pipeline #16794

Uh oh!

Conversation

RkGrit commented Nov 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Key Changes

Reviewed changes

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

SpriCoder left a comment

Choose a reason for hiding this comment

Uh oh!

CRZbulabula left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

CRZbulabula Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

CRZbulabula left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

CRZbulabula Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

RkGrit commented Nov 24, 2025 •

edited

Loading