This is an automated email from the ASF dual-hosted git repository.
CRZbulabula 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 38ea60e04f9 [AINode] Eliminate transformers registered and fix build
bug (#17615)
38ea60e04f9 is described below
commit 38ea60e04f95c71a2f6a64b2cc8b42cd2c770d76
Author: Yongzao <[email protected]>
AuthorDate: Fri May 8 19:24:31 2026 +0800
[AINode] Eliminate transformers registered and fix build bug (#17615)
---
.../iotdb/ainode/it/AINodeModelManageIT.java | 10 -
iotdb-core/ainode/build_binary.py | 237 +--------------------
.../ainode/iotdb/ainode/core/model/model_info.py | 8 -
.../iotdb/ainode/core/model/model_storage.py | 210 +-----------------
4 files changed, 6 insertions(+), 459 deletions(-)
diff --git
a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
index 9e78c8b025c..42dc8db520e 100644
---
a/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
+++
b/integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
@@ -77,11 +77,6 @@ public class AINodeModelManageIT {
registerUserDefinedModel(statement, modelInfo, "file:///data/chronos2");
callInferenceTest(statement, modelInfo);
dropUserDefinedModel(statement, modelInfo.getModelId());
- errorTest(
- statement,
- "create model origin_chronos using uri
\"file:///data/chronos2_origin\"",
- "1505: 't5' is already used by a Transformers config, pick another
name.");
- statement.execute("drop model origin_chronos");
// Test PytorchModelHubMixin model (mantis) in tree.
modelInfo = new FakeModelInfo("user_mantis", "custom_mantis",
"user_defined", "active");
@@ -100,11 +95,6 @@ public class AINodeModelManageIT {
registerUserDefinedModel(statement, modelInfo, "file:///data/chronos2");
forecastTableFunctionTest(statement, modelInfo);
dropUserDefinedModel(statement, modelInfo.getModelId());
- errorTest(
- statement,
- "create model origin_chronos using uri
\"file:///data/chronos2_origin\"",
- "1505: 't5' is already used by a Transformers config, pick another
name.");
- statement.execute("drop model origin_chronos");
// Test PytorchModelHubMixin model (mantis) in table.
modelInfo = new FakeModelInfo("user_mantis", "custom_mantis",
"user_defined", "active");
diff --git a/iotdb-core/ainode/build_binary.py
b/iotdb-core/ainode/build_binary.py
index f3b7fa1cedf..56f19f43cf5 100644
--- a/iotdb-core/ainode/build_binary.py
+++ b/iotdb-core/ainode/build_binary.py
@@ -22,7 +22,6 @@ PyInstaller build script (Python version)
"""
import os
-import shutil
import subprocess
import sys
from pathlib import Path
@@ -122,7 +121,6 @@ def get_venv_env(venv_dir):
Sets VIRTUAL_ENV and prepends the venv's bin/Scripts directory to PATH
so that tools installed in the venv take precedence.
- Also sets POETRY_VIRTUALENVS_PATH to force poetry to use our venv.
Returns:
dict: Environment variables dictionary
@@ -133,10 +131,6 @@ def get_venv_env(venv_dir):
venv_bin = str(venv_dir / ("Scripts" if sys.platform == "win32" else
"bin"))
env["PATH"] = f"{venv_bin}{os.pathsep}{env.get('PATH', '')}"
- # Force poetry to use our virtual environment by setting
POETRY_VIRTUALENVS_PATH
- # This tells poetry where to look for/create virtual environments
- env["POETRY_VIRTUALENVS_PATH"] = str(venv_dir.parent.absolute())
-
return env
@@ -159,10 +153,11 @@ def install_dependencies(venv_python, venv_dir,
script_dir):
venv_env = get_venv_env(venv_dir)
poetry_exe = get_poetry_executable(venv_dir)
- # Configure poetry settings
+ # Configure poetry to NOT create its own virtual environments.
+ # Poetry will use the already-activated venv via the VIRTUAL_ENV
+ # environment variable set in get_venv_env().
print("Configuring poetry settings...")
try:
- # Set poetry to not create venvs in project directory
subprocess.run(
[str(poetry_exe), "config", "virtualenvs.in-project", "false"],
cwd=str(script_dir),
@@ -171,24 +166,8 @@ def install_dependencies(venv_python, venv_dir,
script_dir):
capture_output=True,
text=True,
)
- # Set poetry virtualenvs path to our venv directory's parent
- # This forces poetry to look for/create venvs in the same location as
our venv
- subprocess.run(
- [
- str(poetry_exe),
- "config",
- "virtualenvs.path",
- str(venv_dir.parent.absolute()),
- ],
- cwd=str(script_dir),
- env=venv_env,
- check=True,
- capture_output=True,
- text=True,
- )
- # Ensure poetry can use virtual environments
subprocess.run(
- [str(poetry_exe), "config", "virtualenvs.create", "true"],
+ [str(poetry_exe), "config", "virtualenvs.create", "false"],
cwd=str(script_dir),
env=venv_env,
check=True,
@@ -197,28 +176,8 @@ def install_dependencies(venv_python, venv_dir,
script_dir):
)
except Exception as e:
print(f"Warning: Failed to configure poetry settings: {e}")
- # Continue anyway, as these may not be critical
- # Remove any existing poetry virtual environments for this project
- # This ensures poetry will use our specified virtual environment
- print("Removing any existing poetry virtual environments...")
- remove_result = subprocess.run(
- [str(poetry_exe), "env", "remove", "--all"],
- cwd=str(script_dir),
- env=venv_env,
- check=False, # Don't fail if no venv exists
- capture_output=True,
- text=True,
- )
- if remove_result.stdout:
- print(remove_result.stdout.strip())
- if remove_result.stderr:
- stderr = remove_result.stderr.strip()
- # Ignore "No virtualenv has been activated" error
- if "no virtualenv" not in stderr.lower():
- print(remove_result.stderr.strip())
-
- # Verify the virtual environment Python is valid before configuring poetry
+ # Verify the virtual environment Python is valid
print(f"Verifying virtual environment Python at: {venv_python}")
if not venv_python.exists():
print(f"ERROR: Virtual environment Python not found at: {venv_python}")
@@ -235,190 +194,8 @@ def install_dependencies(venv_python, venv_dir,
script_dir):
sys.exit(1)
print(f" Python version: {python_version_result.stdout.strip()}")
- # Instead of using poetry env use (which creates new venvs), we'll use a
different approach:
- # 1. Create a symlink from poetry's expected venv location to our venv
- # 2. Or, directly use poetry install with VIRTUAL_ENV set (poetry should
detect it)
- #
- # The issue is that poetry env use creates venvs with hash-based names in
its cache.
- # We need to work around this by either:
- # - Creating a symlink from poetry's expected location to our venv
- # - Or bypassing poetry env use entirely and using poetry install directly
-
- # Strategy: Create a symlink from poetry's expected venv location to our
venv
- # Poetry creates venvs with names like:
<project-name>-<hash>-py<python-version>
- # We need to find out what poetry would name our venv, then create a
symlink
-
- print(f"Configuring poetry to use virtual environment at: {venv_dir}")
-
- # Get poetry's expected venv name by checking what it would create
- # First, let's try poetry env use, but catch if it tries to create a new
venv
- result = subprocess.run(
- [str(poetry_exe), "env", "use", str(venv_python)],
- cwd=str(script_dir),
- env=venv_env,
- check=False,
- capture_output=True,
- text=True,
- )
-
- output_text = (result.stdout or "") + (result.stderr or "")
-
- # If poetry is creating a new venv, we need to stop it and use a different
approach
- if (
- "Creating virtualenv" in output_text
- or "Creating virtual environment" in output_text
- or "Using virtualenv:" in output_text
- ):
- print("Poetry is attempting to create/use a new virtual environment.")
- print(
- "Stopping this and using alternative approach: creating symlink to
our venv..."
- )
-
- # Extract the venv path poetry is trying to create/use
- # Look for patterns like "Using virtualenv: /path/to/venv" or
"Creating virtualenv name in /path"
- import re
-
- poetry_venv_path = None
-
- # Try to extract from "Using virtualenv: /path/to/venv"
- using_match = re.search(r"Using virtualenv:\s*([^\s\n]+)", output_text)
- if using_match:
- poetry_venv_path = Path(using_match.group(1))
-
- # If not found, try to extract from "Creating virtualenv name in /path"
- if not poetry_venv_path:
- creating_match = re.search(
- r"Creating virtualenv[^\n]*in\s+([^\s\n]+)", output_text
- )
- if creating_match:
- venv_dir_path = Path(creating_match.group(1))
- # Extract venv name from the output
- name_match = re.search(r"Creating virtualenv\s+([^\s]+)",
output_text)
- if name_match:
- venv_name = name_match.group(1)
- poetry_venv_path = venv_dir_path / venv_name
-
- # If still not found, try to find any path in pypoetry/virtualenvs
- if not poetry_venv_path:
- pypoetry_match = re.search(
- r"([^\s]+pypoetry[^\s]*virtualenvs[^\s]+)", output_text
- )
- if pypoetry_match:
- poetry_venv_path = Path(pypoetry_match.group(1))
-
- if poetry_venv_path:
- print(f"Poetry wants to create/use venv at: {poetry_venv_path}")
-
- # Remove the venv poetry just created (if it exists)
- if poetry_venv_path.exists() and poetry_venv_path.is_dir():
- print(f"Removing poetry's newly created venv:
{poetry_venv_path}")
- shutil.rmtree(poetry_venv_path, ignore_errors=True)
-
- # Create a symlink from poetry's expected location to our venv
- print(f"Creating symlink from {poetry_venv_path} to {venv_dir}")
- try:
- if poetry_venv_path.exists() or poetry_venv_path.is_symlink():
- if poetry_venv_path.is_symlink():
- poetry_venv_path.unlink()
- elif poetry_venv_path.is_dir():
- shutil.rmtree(poetry_venv_path, ignore_errors=True)
- poetry_venv_path.parent.mkdir(parents=True, exist_ok=True)
- poetry_venv_path.symlink_to(venv_dir)
- print(f"Symlink created successfully")
- except Exception as e:
- print(f"WARNING: Failed to create symlink: {e}")
- print("Will try to use poetry install directly with
VIRTUAL_ENV set")
- else:
- print("Could not determine poetry's venv path from output")
- print(f"Output was: {output_text}")
- else:
- if result.stdout:
- print(result.stdout.strip())
- if result.stderr:
- stderr = result.stderr.strip()
- if stderr:
- print(f"Poetry output: {stderr}")
-
- # Verify poetry is using the correct virtual environment BEFORE running
lock/install
- # This is critical - if poetry uses the wrong venv, dependencies won't be
installed correctly
- print("Verifying poetry virtual environment...")
-
- # Wait a moment for symlink to be recognized (if we created one)
- import time
-
- time.sleep(0.5)
-
- verify_result = subprocess.run(
- [str(poetry_exe), "env", "info", "--path"],
- cwd=str(script_dir),
- env=venv_env,
- check=False, # Don't fail if poetry hasn't activated a venv yet
- capture_output=True,
- text=True,
- )
-
- expected_venv_path_resolved = str(Path(venv_dir.absolute()).resolve())
-
- # If poetry env info fails, it might mean poetry hasn't activated the venv
yet
- if verify_result.returncode != 0:
- print(
- "Warning: poetry env info failed, poetry may not have activated
the virtual environment yet"
- )
- print(
- "This may be okay if we created a symlink - poetry should use it
when running commands"
- )
- poetry_venv_path_resolved = None
- else:
- poetry_venv_path = verify_result.stdout.strip()
-
- # Normalize paths for comparison (resolve symlinks, etc.)
- poetry_venv_path_resolved = str(Path(poetry_venv_path).resolve())
-
- # Only verify path if we successfully got poetry's venv path
- if poetry_venv_path_resolved is not None:
- if poetry_venv_path_resolved != expected_venv_path_resolved:
- print(
- f"ERROR: Poetry is using {poetry_venv_path}, but expected
{expected_venv_path_resolved}"
- )
- print(
- "Poetry must use the virtual environment we created for the
build to work correctly."
- )
- print("The symlink approach may not have worked. Please check the
symlink.")
- sys.exit(1)
- else:
- print(f"Poetry is correctly using virtual environment:
{poetry_venv_path}")
- else:
- print("Warning: Could not verify poetry virtual environment path")
- print(
- "Continuing anyway - poetry should use the venv via symlink or
VIRTUAL_ENV"
- )
-
# Update lock file and install dependencies
- # Re-verify environment before each command to ensure poetry doesn't
switch venvs
- def verify_poetry_env():
- verify_result = subprocess.run(
- [str(poetry_exe), "env", "info", "--path"],
- cwd=str(script_dir),
- env=venv_env,
- check=False, # Don't fail if poetry env info is not available
- capture_output=True,
- text=True,
- )
- if verify_result.returncode == 0:
- current_path = str(Path(verify_result.stdout.strip()).resolve())
- expected_path = str(Path(venv_dir.absolute()).resolve())
- if current_path != expected_path:
- print(
- f"ERROR: Poetry switched to different virtual environment:
{current_path}"
- )
- print(f"Expected: {expected_path}")
- sys.exit(1)
- # If poetry env info fails, we can't verify, but continue anyway
- # Poetry should still use the Python we specified via env use
- return True
-
print("Running poetry lock...")
- verify_poetry_env() # Verify before lock
result = subprocess.run(
[str(poetry_exe), "lock"],
cwd=str(script_dir),
@@ -434,7 +211,6 @@ def install_dependencies(venv_python, venv_dir, script_dir):
if result.returncode != 0:
print(f"ERROR: poetry lock failed with exit code {result.returncode}")
sys.exit(1)
- verify_poetry_env() # Verify after lock
accelerator = detect_accelerator()
print(f"Selected accelerator: {accelerator}")
@@ -447,12 +223,9 @@ def install_dependencies(venv_python, venv_dir,
script_dir):
check=True,
text=True,
)
- verify_poetry_env() # Verify before install
poetry_install_with_accel(poetry_exe, script_dir, venv_env, accelerator)
- verify_poetry_env() # Verify after install
# Verify installation by checking if key packages are installed
- # This is critical - if packages aren't installed, PyInstaller won't find
them
print("Verifying package installation...")
test_packages = ["torch", "transformers", "tokenizers"]
missing_packages = []
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py
b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py
index 0c63b356e7a..50b0aaf66bc 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_info.py
@@ -32,7 +32,6 @@ class ModelInfo:
repo_id: str = "",
auto_map: Optional[Dict] = None,
hub_mixin_cls: Optional[str] = None,
- transformers_registered: bool = False,
):
self.model_id = model_id
self.model_type = model_type
@@ -42,7 +41,6 @@ class ModelInfo:
self.repo_id = repo_id
self.auto_map = auto_map
self.hub_mixin_cls = hub_mixin_cls
- self.transformers_registered = transformers_registered
def __repr__(self):
return (
@@ -119,7 +117,6 @@ BUILTIN_HF_TRANSFORMERS_MODEL_MAP = {
"AutoConfig": "configuration_timer.TimerConfig",
"AutoModelForCausalLM": "modeling_timer.TimerForPrediction",
},
- transformers_registered=True,
),
"sundial": ModelInfo(
model_id="sundial",
@@ -132,7 +129,6 @@ BUILTIN_HF_TRANSFORMERS_MODEL_MAP = {
"AutoConfig": "configuration_sundial.SundialConfig",
"AutoModelForCausalLM": "modeling_sundial.SundialForPrediction",
},
- transformers_registered=True,
),
"chronos2": ModelInfo(
model_id="chronos2",
@@ -145,7 +141,6 @@ BUILTIN_HF_TRANSFORMERS_MODEL_MAP = {
"AutoConfig": "config.Chronos2CoreConfig",
"AutoModelForCausalLM": "model.Chronos2Model",
},
- transformers_registered=True,
),
"moirai2": ModelInfo(
model_id="moirai2",
@@ -158,7 +153,6 @@ BUILTIN_HF_TRANSFORMERS_MODEL_MAP = {
"AutoConfig": "configuration_moirai2.Moirai2Config",
"AutoModelForCausalLM": "modeling_moirai2.Moirai2ForPrediction",
},
- transformers_registered=True,
),
"toto": ModelInfo(
model_id="toto",
@@ -171,7 +165,6 @@ BUILTIN_HF_TRANSFORMERS_MODEL_MAP = {
"AutoConfig": "configuration_toto.TotoConfig",
"AutoModelForCausalLM": "modeling_toto.TotoForPrediction",
},
- transformers_registered=True,
),
"moment": ModelInfo(
model_id="moment",
@@ -184,6 +177,5 @@ BUILTIN_HF_TRANSFORMERS_MODEL_MAP = {
"AutoConfig": "configuration_moment.MomentConfig",
"AutoModelForCausalLM": "modeling_moment.MomentForPrediction",
},
- transformers_registered=True,
),
}
diff --git a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
index b7799df67a5..b824fe84fc1 100644
--- a/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
+++ b/iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py
@@ -20,16 +20,9 @@ import concurrent.futures
import json
import os
import shutil
-from pathlib import Path
from typing import Dict, Optional
-from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
-from transformers import (
- AutoConfig,
- AutoModelForCausalLM,
- PretrainedConfig,
- PreTrainedModel,
-)
+from huggingface_hub import hf_hub_download
from iotdb.ainode.core.config import AINodeDescriptor
from iotdb.ainode.core.constant import TSStatusCode
@@ -56,10 +49,8 @@ from iotdb.ainode.core.model.utils import (
_fetch_model_from_local,
ensure_init_file,
get_parsed_uri,
- import_class_from_path,
load_model_config_in_json,
parse_uri_type,
- temporary_sys_path,
validate_model_files,
)
from iotdb.ainode.core.util.lock import ModelLockPool
@@ -242,16 +233,9 @@ class ModelStorage:
pipeline_cls=pipeline_cls,
auto_map=auto_map,
hub_mixin_cls=hub_mixin_cls,
- transformers_registered=False, # Lazy registration
)
with self._lock_pool.get_lock(model_id).write_lock():
self._models[ModelCategory.USER_DEFINED.value][model_id] =
model_info
- if self.ensure_transformers_registered(model_id) is None:
- model_info.state = ModelStates.INACTIVE
- else:
- model_info.transformers_registered = True
- with self._lock_pool.get_lock(model_id).write_lock():
- self._models[ModelCategory.USER_DEFINED.value][model_id] =
model_info
# ==================== Registration Methods ====================
@@ -267,7 +251,6 @@ class ModelStorage:
Raises:
ModelExistedException: If the model_id already exists.
InvalidModelUriException: If the URI format is invalid.
- Exception: For other errors during transformers model registration.
"""
if self.is_model_registered(model_id):
@@ -303,198 +286,11 @@ class ModelStorage:
pipeline_cls=pipeline_cls,
auto_map=auto_map,
hub_mixin_cls=hub_mixin_cls,
- transformers_registered=False, # Register later
)
self._models[ModelCategory.USER_DEFINED.value][model_id] =
model_info
- if auto_map:
- # Transformers model: immediately register to Transformers
autoloading mechanism
- try:
- if self._register_transformers_model(model_info):
- model_info.transformers_registered = True
- except Exception as e:
- model_info.state = ModelStates.INACTIVE
- logger.error(
- f"Failed to register Transformers model {model_id},
because {e}"
- )
- raise e
- elif hub_mixin_cls:
- # PyTorchModelHubMixin model: immediately register
- try:
- if self._register_hub_mixin_model(model_info):
- model_info.transformers_registered = True
- except Exception as e:
- model_info.state = ModelStates.INACTIVE
- logger.error(
- f"Failed to register HubMixin model {model_id},
because {e}"
- )
- raise e
- else:
- # Other type models: only log
- self._register_other_model(model_info)
-
logger.info(f"Successfully registered model {model_id} from URI:
{uri}")
- def _register_transformers_model(self, model_info: ModelInfo) -> bool:
- """
- Register Transformers model to autoloading mechanism (internal method)
- Returns:
- True if registration is successful
- Raises:
- Exception: Transformers internal exception if registration fails
- ValueError: If class is invalid
- """
- auto_map = model_info.auto_map
- if not auto_map:
- return False
-
- auto_config_path = auto_map.get("AutoConfig")
- auto_model_path = auto_map.get("AutoModelForCausalLM")
-
- try:
- model_path = os.path.join(
- self._models_dir, model_info.category.value,
model_info.model_id
- )
- module_parent = str(Path(model_path).parent.absolute())
- with temporary_sys_path(module_parent):
- config_class = import_class_from_path(
- model_info.model_id, auto_config_path
- )
- # Validate config_class is a subclass of PretrainedConfig
- if not (
- isinstance(config_class, type)
- and issubclass(config_class, PretrainedConfig)
- ):
- raise ValueError(
- f"AutoConfig class '{auto_config_path}' must be a
subclass of PretrainedConfig"
- )
- AutoConfig.register(model_info.model_type, config_class)
- logger.info(
- f"Registered AutoConfig: {model_info.model_type} ->
{auto_config_path}"
- )
-
- model_class = import_class_from_path(
- model_info.model_id, auto_model_path
- )
- # Validate model_class is a subclass of PreTrainedModel
- if not (
- isinstance(model_class, type)
- and issubclass(model_class, PreTrainedModel)
- ):
- raise ValueError(
- f"AutoModelForCausalLM class '{auto_model_path}' must
be a subclass of PreTrainedModel"
- )
- AutoModelForCausalLM.register(config_class, model_class)
- logger.info(
- f"Registered AutoModelForCausalLM: {config_class.__name__}
-> {auto_model_path}"
- )
- return True
- except Exception as e:
- logger.warning(
- f"Failed to register Transformers model {model_info.model_id}:
{e}. Model may still work via auto_map, but ensure module path is correct."
- )
- raise e
-
- def _register_hub_mixin_model(self, model_info: ModelInfo) -> bool:
- """
- Register PyTorchModelHubMixin model (internal method).
- For now, just validate the class.
-
- Returns:
- True if registration is successful
- Raises:
- ValueError: If class is invalid
- Exception: For other errors
- """
- hub_mixin_cls = model_info.hub_mixin_cls
- if not hub_mixin_cls:
- return False
-
- try:
- model_path = os.path.join(
- self._models_dir, model_info.category.value,
model_info.model_id
- )
- module_parent = str(Path(model_path).parent.absolute())
- with temporary_sys_path(module_parent):
- model_class = import_class_from_path(model_info.model_id,
hub_mixin_cls)
-
- # Validate that the class inherits from PyTorchModelHubMixin
- if not issubclass(model_class, PyTorchModelHubMixin):
- raise ValueError(
- f"Class '{model_class}' does not inherit from "
- "PyTorchModelHubMixin."
- )
-
- logger.info(
- f"Registered PyTorchModelHubMixin model: "
- f"{model_info.model_id} -> {hub_mixin_cls}"
- )
- return True
-
- except Exception as e:
- logger.warning(
- f"Failed to register PyTorchModelHubMixin model
{model_info.model_id}: {e}."
- )
- raise e
-
- def _register_other_model(self, model_info: ModelInfo):
- """Register other type models (non-Transformers models)"""
- logger.info(
- f"Registered other type model: {model_info.model_id}
({model_info.model_type})"
- )
-
- def ensure_transformers_registered(self, model_id: str) -> ModelInfo |
None:
- """
- Ensure Transformers model is registered.
- Returns:
- ModelInfo | None: None if registration failed, otherwise returns
the corresponding ModelInfo
- """
- # Use lock to protect entire check-execute process
- with self._lock_pool.get_lock(model_id).write_lock():
- # Directly access _models dictionary (avoid calling get_model_info
which may cause deadlock)
- model_info = None
- for category_dict in self._models.values():
- if model_id in category_dict:
- model_info = category_dict[model_id]
- break
-
- if not model_info:
- return None
-
- # If already registered, return directly
- if model_info.transformers_registered:
- return model_info
-
- # If no auto_map, not a Transformers model, mark as registered
(avoid duplicate checks)
- if (
- not model_info.auto_map
- or model_id in BUILTIN_HF_TRANSFORMERS_MODEL_MAP.keys()
- ):
- model_info.transformers_registered = True
- return model_info
-
- # Execute registration (under lock protection)
- try:
- if self._register_transformers_model(model_info):
- model_info.transformers_registered = True
- logger.info(
- f"Model {model_id} successfully registered to
Transformers"
- )
- return model_info
- else:
- model_info.state = ModelStates.INACTIVE
- logger.error(f"Model {model_id} failed to register to
Transformers")
- return None
-
- except Exception as e:
- # Ensure state consistency in exception cases
- model_info.state = ModelStates.INACTIVE
- model_info.transformers_registered = False
- logger.error(
- f"Exception occurred while registering model {model_id} to
Transformers: {e}"
- )
- return None
-
# ==================== Show and Delete Models ====================
def show_models(self, req: TShowModelsReq) -> TShowModelsResp:
@@ -614,10 +410,6 @@ class ModelStorage:
def is_model_registered(self, model_id: str) -> bool:
"""Check if model is registered (search in _models)"""
- # Lazy registration: if it's a Transformers model and not registered,
register it first
- if self.ensure_transformers_registered(model_id) is None:
- return False
-
with self._lock_pool.get_lock("").read_lock():
for category_dict in self._models.values():
if model_id in category_dict:
