Skip to content

Commit 49c625b

Browse files
authored
[AINode] Package AINode via PyInstaller (#16707)
1 parent 43b89b9 commit 49c625b

File tree

22 files changed

+1028
-686
lines changed

22 files changed

+1028
-686
lines changed

.github/workflows/cluster-it-1c1d1a.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,5 +59,5 @@ jobs:
5959
uses: actions/upload-artifact@v4
6060
with:
6161
name: cluster-log-ainode-${{ matrix.os }}
62-
path: integration-test/target/ainode-logs
62+
path: integration-test/target/*-logs
6363
retention-days: 30

integration-test/src/assembly/mpp-test.xml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@
6363
</fileSet>
6464
<fileSet>
6565
<outputDirectory>lib</outputDirectory>
66-
<directory>${project.basedir}/../iotdb-core/ainode/dist/</directory>
66+
<directory>${project.basedir}/../iotdb-core/ainode/dist/ainode/</directory>
6767
<fileMode>0755</fileMode>
6868
</fileSet>
6969
</fileSets>

integration-test/src/main/java/org/apache/iotdb/it/env/cluster/node/AINodeWrapper.java

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,9 @@
2525
import org.apache.tsfile.external.commons.io.file.PathUtils;
2626
import org.slf4j.Logger;
2727

28-
import java.io.BufferedWriter;
2928
import java.io.File;
30-
import java.io.FileWriter;
29+
import java.io.FileInputStream;
30+
import java.io.FileOutputStream;
3131
import java.io.IOException;
3232
import java.nio.file.Files;
3333
import java.nio.file.LinkOption;
@@ -37,6 +37,7 @@
3737
import java.nio.file.StandardCopyOption;
3838
import java.util.ArrayList;
3939
import java.util.List;
40+
import java.util.Properties;
4041
import java.util.stream.Stream;
4142

4243
import static org.apache.iotdb.it.env.cluster.ClusterConstant.AI_NODE_NAME;
@@ -62,15 +63,19 @@ public class AINodeWrapper extends AbstractNodeWrapper {
6263
public static final String CACHE_BUILT_IN_MODEL_PATH = "/data/ainode/models/weights";
6364

6465
private void replaceAttribute(String[] keys, String[] values, String filePath) {
65-
try (BufferedWriter writer = new BufferedWriter(new FileWriter(filePath, true))) {
66-
for (int i = 0; i < keys.length; i++) {
67-
String line = keys[i] + "=" + values[i];
68-
writer.newLine();
69-
writer.write(line);
70-
}
66+
Properties props = new Properties();
67+
try (FileInputStream in = new FileInputStream(filePath)) {
68+
props.load(in);
69+
} catch (IOException e) {
70+
logger.warn("Failed to load existing AINode properties from {}, because: ", filePath, e);
71+
}
72+
for (int i = 0; i < keys.length; i++) {
73+
props.setProperty(keys[i], values[i]);
74+
}
75+
try (FileOutputStream out = new FileOutputStream(filePath)) {
76+
props.store(out, "Updated by AINode integration-test env");
7177
} catch (IOException e) {
72-
logger.error(
73-
"Failed to set attribute for AINode in file: {} because {}", filePath, e.getMessage());
78+
logger.error("Failed to save properties to {}, because:", filePath, e);
7479
}
7580
}
7681

integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeConcurrentInferenceIT.java

Lines changed: 1 addition & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -90,33 +90,6 @@ private static void prepareDataForTableModel() throws SQLException {
9090
}
9191
}
9292

93-
// @Test
94-
public void concurrentCPUCallInferenceTest() throws SQLException, InterruptedException {
95-
concurrentCPUCallInferenceTest("timer_xl");
96-
concurrentCPUCallInferenceTest("sundial");
97-
}
98-
99-
private void concurrentCPUCallInferenceTest(String modelId)
100-
throws SQLException, InterruptedException {
101-
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
102-
Statement statement = connection.createStatement()) {
103-
final int threadCnt = 4;
104-
final int loop = 10;
105-
final int predictLength = 96;
106-
statement.execute(String.format("LOAD MODEL %s TO DEVICES 'cpu'", modelId));
107-
checkModelOnSpecifiedDevice(statement, modelId, "cpu");
108-
concurrentInference(
109-
statement,
110-
String.format(
111-
"CALL INFERENCE(%s, 'SELECT s FROM root.AI', predict_length=%d)",
112-
modelId, predictLength),
113-
threadCnt,
114-
loop,
115-
predictLength);
116-
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES 'cpu'", modelId));
117-
}
118-
}
119-
12093
// @Test
12194
public void concurrentGPUCallInferenceTest() throws SQLException, InterruptedException {
12295
concurrentGPUCallInferenceTest("timer_xl");
@@ -150,39 +123,6 @@ private void concurrentGPUCallInferenceTest(String modelId)
150123
String forecastUDTFSql =
151124
"SELECT forecast(s, 'MODEL_ID'='%s', 'PREDICT_LENGTH'='%d') FROM root.AI";
152125

153-
@Test
154-
public void concurrentCPUForecastTest() throws SQLException, InterruptedException {
155-
concurrentCPUForecastTest("timer_xl", forecastUDTFSql);
156-
concurrentCPUForecastTest("sundial", forecastUDTFSql);
157-
concurrentCPUForecastTest("timer_xl", forecastTableFunctionSql);
158-
concurrentCPUForecastTest("sundial", forecastTableFunctionSql);
159-
}
160-
161-
private void concurrentCPUForecastTest(String modelId, String selectSQL)
162-
throws SQLException, InterruptedException {
163-
try (Connection connection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
164-
Statement statement = connection.createStatement()) {
165-
final int threadCnt = 4;
166-
final int loop = 10;
167-
final int predictLength = 96;
168-
statement.execute(String.format("LOAD MODEL %s TO DEVICES 'cpu'", modelId));
169-
checkModelOnSpecifiedDevice(statement, modelId, "cpu");
170-
long startTime = System.currentTimeMillis();
171-
concurrentInference(
172-
statement,
173-
String.format(selectSQL, modelId, predictLength),
174-
threadCnt,
175-
loop,
176-
predictLength);
177-
long endTime = System.currentTimeMillis();
178-
LOGGER.info(
179-
String.format(
180-
"Model %s concurrent inference %d reqs (%d threads, %d loops) in CPU takes time: %dms",
181-
modelId, threadCnt * loop, threadCnt, loop, endTime - startTime));
182-
statement.execute(String.format("UNLOAD MODEL %s FROM DEVICES 'cpu'", modelId));
183-
}
184-
}
185-
186126
@Test
187127
public void concurrentGPUForecastTest() throws SQLException, InterruptedException {
188128
concurrentGPUForecastTest("timer_xl", forecastUDTFSql);
@@ -221,7 +161,7 @@ private void checkModelOnSpecifiedDevice(Statement statement, String modelId, St
221161
throws SQLException, InterruptedException {
222162
Set<String> targetDevices = ImmutableSet.copyOf(device.split(","));
223163
LOGGER.info("Checking model: {} on target devices: {}", modelId, targetDevices);
224-
for (int retry = 0; retry < 20; retry++) {
164+
for (int retry = 0; retry < 200; retry++) {
225165
Set<String> foundDevices = new HashSet<>();
226166
try (final ResultSet resultSet =
227167
statement.executeQuery(String.format("SHOW LOADED MODELS '%s'", device))) {

iotdb-core/ainode/.gitignore

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,6 @@
1414
# generated by maven
1515
/iotdb/ainode/conf/
1616

17-
# .whl of ainode, generated by Poetry
17+
# generated by pyinstaller
1818
/dist/
19-
20-
# the config to build ainode, it will be generated automatically
21-
pyproject.toml
19+
/build/

iotdb-core/ainode/ainode.spec

Lines changed: 199 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,199 @@
1+
# -*- mode: python ; coding: utf-8 -*-
2+
#
3+
# Licensed to the Apache Software Foundation (ASF) under one
4+
# or more contributor license agreements. See the NOTICE file
5+
# distributed with this work for additional information
6+
# regarding copyright ownership. The ASF licenses this file
7+
# to you under the Apache License, Version 2.0 (the
8+
# "License"); you may not use this file except in compliance
9+
# with the License. You may obtain a copy of the License at
10+
#
11+
# http://www.apache.org/licenses/LICENSE-2.0
12+
#
13+
# Unless required by applicable law or agreed to in writing,
14+
# software distributed under this License is distributed on an
15+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16+
# KIND, either express or implied. See the License for the
17+
# specific language governing permissions and limitations
18+
# under the License.
19+
#
20+
21+
from pathlib import Path
22+
23+
# Get project root directory
24+
project_root = Path(SPECPATH).parent
25+
26+
block_cipher = None
27+
28+
# Auto-collect all submodules of large dependency libraries
29+
# Using collect_all automatically includes all dependencies and avoids manual maintenance of hiddenimports
30+
from PyInstaller.utils.hooks import collect_all, collect_submodules, collect_data_files
31+
32+
# Collect only essential data files and binaries for large libraries
33+
# Using collect_all for all submodules slows down startup significantly.
34+
# However, for certain libraries with many dynamic imports (e.g., torch, transformers, safetensors),
35+
# collect_all is necessary to ensure all required modules are included.
36+
# For other libraries, we use lighter-weight collection methods to improve startup time.
37+
all_datas = []
38+
all_binaries = []
39+
all_hiddenimports = []
40+
41+
# Only collect essential data files and binaries for critical libraries
42+
# This reduces startup time by avoiding unnecessary module imports
43+
essential_libraries = {
44+
'torch': True, # Keep collect_all for torch as it has many dynamic imports
45+
'transformers': True, # Keep collect_all for transformers
46+
'safetensors': True, # Keep collect_all for safetensors
47+
}
48+
49+
# For other libraries, use selective collection to speed up startup
50+
other_libraries = ['sktime', 'scipy', 'pandas', 'sklearn', 'statsmodels', 'optuna']
51+
52+
for lib in essential_libraries:
53+
try:
54+
lib_datas, lib_binaries, lib_hiddenimports = collect_all(lib)
55+
all_datas.extend(lib_datas)
56+
all_binaries.extend(lib_binaries)
57+
all_hiddenimports.extend(lib_hiddenimports)
58+
except Exception:
59+
pass
60+
61+
# For other libraries, only collect submodules (lighter weight)
62+
# This relies on PyInstaller's dependency analysis to include what's actually used
63+
for lib in other_libraries:
64+
try:
65+
submodules = collect_submodules(lib)
66+
all_hiddenimports.extend(submodules)
67+
# Only collect essential data files and binaries, not all submodules
68+
# This significantly reduces startup time
69+
try:
70+
lib_datas, lib_binaries, _ = collect_all(lib)
71+
all_datas.extend(lib_datas)
72+
all_binaries.extend(lib_binaries)
73+
except Exception:
74+
# If collect_all fails, try collect_data_files for essential data only
75+
try:
76+
lib_datas = collect_data_files(lib)
77+
all_datas.extend(lib_datas)
78+
except Exception:
79+
pass
80+
except Exception:
81+
pass
82+
83+
# Project-specific packages that need their submodules collected
84+
# Only list top-level packages - collect_submodules will recursively collect all submodules
85+
TOP_LEVEL_PACKAGES = [
86+
'iotdb.ainode.core', # This will include all sub-packages: manager, model, inference, etc.
87+
'iotdb.thrift', # This will include all thrift sub-packages
88+
]
89+
90+
# Collect all submodules for project packages automatically
91+
# Using top-level packages avoids duplicate collection
92+
for package in TOP_LEVEL_PACKAGES:
93+
try:
94+
submodules = collect_submodules(package)
95+
all_hiddenimports.extend(submodules)
96+
except Exception:
97+
# If package doesn't exist or collection fails, add the package itself
98+
all_hiddenimports.append(package)
99+
100+
# Add parent packages to ensure they are included
101+
all_hiddenimports.extend(['iotdb', 'iotdb.ainode'])
102+
103+
# Multiprocessing support for PyInstaller
104+
# When using multiprocessing with PyInstaller, we need to ensure proper handling
105+
multiprocessing_modules = [
106+
'multiprocessing',
107+
'multiprocessing.spawn',
108+
'multiprocessing.popen_spawn_posix',
109+
'multiprocessing.popen_spawn_win32',
110+
'multiprocessing.popen_fork',
111+
'multiprocessing.popen_forkserver',
112+
'multiprocessing.context',
113+
'multiprocessing.reduction',
114+
'multiprocessing.util',
115+
'torch.multiprocessing',
116+
'torch.multiprocessing.spawn',
117+
]
118+
119+
# Additional dependencies that may need explicit import
120+
# These are external libraries that might use dynamic imports
121+
external_dependencies = [
122+
'huggingface_hub',
123+
'tokenizers',
124+
'hf_xet',
125+
'einops',
126+
'dynaconf',
127+
'tzlocal',
128+
'thrift',
129+
'psutil',
130+
'requests',
131+
]
132+
133+
all_hiddenimports.extend(multiprocessing_modules)
134+
all_hiddenimports.extend(external_dependencies)
135+
136+
# Analyze main entry file
137+
# Note: Do NOT add virtual environment site-packages to pathex manually.
138+
# When PyInstaller is run from the virtual environment's Python, it automatically
139+
# detects and uses the virtual environment's site-packages.
140+
a = Analysis(
141+
['iotdb/ainode/core/script.py'],
142+
pathex=[str(project_root)],
143+
binaries=all_binaries,
144+
datas=all_datas,
145+
hiddenimports=all_hiddenimports,
146+
hookspath=[],
147+
hooksconfig={},
148+
runtime_hooks=[],
149+
excludes=[
150+
# Exclude unnecessary modules to reduce size and improve startup time
151+
# Note: Do not exclude unittest, as torch and other libraries require it
152+
# Only exclude modules that are definitely not used and not required by dependencies
153+
'matplotlib',
154+
'IPython',
155+
'jupyter',
156+
'notebook',
157+
'pytest',
158+
'test',
159+
'tests'
160+
],
161+
win_no_prefer_redirects=False,
162+
win_private_assemblies=False,
163+
cipher=block_cipher,
164+
noarchive=True, # Set to True to speed up startup - files are not archived into PYZ
165+
)
166+
167+
# Package all PYZ files
168+
pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher)
169+
170+
# Create executable (onedir mode for faster startup)
171+
exe = EXE(
172+
pyz,
173+
a.scripts,
174+
[],
175+
exclude_binaries=True,
176+
name='ainode',
177+
debug=False,
178+
bootloader_ignore_signals=False,
179+
strip=False,
180+
upx=True,
181+
console=True,
182+
disable_windowed_traceback=False,
183+
argv_emulation=False,
184+
target_arch=None,
185+
codesign_identity=None,
186+
entitlements_file=None,
187+
)
188+
189+
# Collect all files into a directory (onedir mode)
190+
coll = COLLECT(
191+
exe,
192+
a.binaries,
193+
a.zipfiles,
194+
a.datas,
195+
strip=False,
196+
upx=True,
197+
upx_exclude=[],
198+
name='ainode',
199+
)

iotdb-core/ainode/ainode.xml

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@
4242
</file>
4343
</files>
4444
<fileSets>
45+
<fileSet>
46+
<directory>iotdb/ainode/conf</directory>
47+
<outputDirectory>conf</outputDirectory>
48+
</fileSet>
4549
<fileSet>
4650
<directory>resources/conf</directory>
4751
<outputDirectory>conf</outputDirectory>
@@ -52,19 +56,8 @@
5256
<fileMode>0755</fileMode>
5357
</fileSet>
5458
<fileSet>
55-
<directory>dist</directory>
59+
<directory>dist/ainode</directory>
5660
<outputDirectory>lib</outputDirectory>
57-
<includes>
58-
<include>*.whl</include>
59-
</includes>
60-
</fileSet>
61-
<fileSet>
62-
<directory>${project.basedir}/../../scripts/conf</directory>
63-
<outputDirectory>conf</outputDirectory>
64-
<includes>
65-
<include>ainode-env.*</include>
66-
<include>**/ainode-env.*</include>
67-
</includes>
6861
<fileMode>0755</fileMode>
6962
</fileSet>
7063
<fileSet>

0 commit comments

Comments
 (0)