Skip to content

[AINode] Eliminate transformers registered and fix build bug #17615

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
CRZbulabula merged 2 commits into from
May 8, 2026
+6 −459
Merged

[AINode] Eliminate transformers registered and fix build bug #17615

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
908de7f
f
CRZbulabula May 7, 2026
c1f9996
remove useless test
CRZbulabula May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 0 additions & 10 deletions integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeModelManageIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,6 @@ public void userDefinedModelManagementTestInTree() throws SQLException, Interrup
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");
Expand All @@ -100,11 +95,6 @@ public void userDefinedModelManagementTestInTable() throws SQLException, Interru
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");
Expand Down
237 changes: 5 additions & 232 deletions iotdb-core/ainode/build_binary.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
"""

import os
import shutil
import subprocess
import sys
from pathlib import Path
Expand Down Expand Up @@ -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
Expand All @@ -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


Expand All @@ -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),
Expand All @@ -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,
Expand All @@ -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}")
Expand All @@ -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),
Expand All @@ -434,7 +211,6 @@ def verify_poetry_env():
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}")
Expand All @@ -447,12 +223,9 @@ def verify_poetry_env():
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 = []
Expand Down
8 changes: 0 additions & 8 deletions iotdb-core/ainode/iotdb/ainode/core/model/model_info.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@ def __init__(
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
Expand All @@ -42,7 +41,6 @@ def __init__(
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 (
Expand Down Expand Up @@ -119,7 +117,6 @@ def __repr__(self):
"AutoConfig": "configuration_timer.TimerConfig",
"AutoModelForCausalLM": "modeling_timer.TimerForPrediction",
},
transformers_registered=True,
),
"sundial": ModelInfo(
model_id="sundial",
Expand All @@ -132,7 +129,6 @@ def __repr__(self):
"AutoConfig": "configuration_sundial.SundialConfig",
"AutoModelForCausalLM": "modeling_sundial.SundialForPrediction",
},
transformers_registered=True,
),
"chronos2": ModelInfo(
model_id="chronos2",
Expand All @@ -145,7 +141,6 @@ def __repr__(self):
"AutoConfig": "config.Chronos2CoreConfig",
"AutoModelForCausalLM": "model.Chronos2Model",
},
transformers_registered=True,
),
"moirai2": ModelInfo(
model_id="moirai2",
Expand All @@ -158,7 +153,6 @@ def __repr__(self):
"AutoConfig": "configuration_moirai2.Moirai2Config",
"AutoModelForCausalLM": "modeling_moirai2.Moirai2ForPrediction",
},
transformers_registered=True,
),
"toto": ModelInfo(
model_id="toto",
Expand All @@ -171,7 +165,6 @@ def __repr__(self):
"AutoConfig": "configuration_toto.TotoConfig",
"AutoModelForCausalLM": "modeling_toto.TotoForPrediction",
},
transformers_registered=True,
),
"moment": ModelInfo(
model_id="moment",
Expand All @@ -184,6 +177,5 @@ def __repr__(self):
"AutoConfig": "configuration_moment.MomentConfig",
"AutoModelForCausalLM": "modeling_moment.MomentForPrediction",
},
transformers_registered=True,
),
}
Loading
Loading