This is an automated email from the ASF dual-hosted git repository. yongzao pushed a commit to branch update-build-process in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 00ae52492c5546fa986ccc13c03729a46adb7cba Author: Yongzao <[email protected]> AuthorDate: Tue Jan 13 17:28:43 2026 +0800 finish --- iotdb-core/ainode/build_binary.py | 54 +++++++++++++++------- .../ainode/iotdb/ainode/core/manager/utils.py | 24 ---------- 2 files changed, 38 insertions(+), 40 deletions(-) diff --git a/iotdb-core/ainode/build_binary.py b/iotdb-core/ainode/build_binary.py index 01b0ef67e8f..c943de41581 100644 --- a/iotdb-core/ainode/build_binary.py +++ b/iotdb-core/ainode/build_binary.py @@ -31,7 +31,7 @@ from pathlib import Path def get_venv_base_dir(): """ Get the base directory for virtual environments outside the project. - + Returns: Path: Base directory path - Linux/macOS: ~/.cache/iotdb-ainode-build/ @@ -51,13 +51,13 @@ def get_venv_base_dir(): def setup_venv(): """ Create virtual environment outside the project directory. - + The virtual environment is created in a platform-specific location: - Linux/macOS: ~/.cache/iotdb-ainode-build/<project-name>/ - Windows: %LOCALAPPDATA%\\iotdb-ainode-build\\<project-name>\\ - + The same venv is reused across multiple builds of the same project. - + Returns: Path: Path to the virtual environment directory """ @@ -433,6 +433,9 @@ def install_dependencies(venv_python, venv_dir, script_dir): print(result.stderr) verify_poetry_env() # Verify after lock + accelerator = detect_accelerator() + print(f"Selected accelerator: {accelerator}") + print("Running poetry install...") subprocess.run( [str(poetry_exe), "lock"], @@ -443,18 +446,7 @@ def install_dependencies(venv_python, venv_dir, script_dir): text=True, ) verify_poetry_env() # Verify before install - result = subprocess.run( - [str(poetry_exe), "install"], - cwd=str(script_dir), - env=venv_env, - check=True, - capture_output=True, - text=True, - ) - if result.stdout: - print(result.stdout) - if result.stderr: - print(result.stderr) + 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 @@ -520,6 +512,36 @@ def check_pyinstaller(venv_python): return False +def detect_accelerator(): + """Auto-detect accelerator: prefer NPU if available, else CUDA GPU, otherwise CPU.""" + + # Try NVIDIA CUDA detection + try: + cuda_result = subprocess.run( + ["nvidia-smi", "-L"], capture_output=True, text=True, check=False + ) + if cuda_result.returncode == 0 and "GPU" in cuda_result.stdout: + return "cuda" + except FileNotFoundError: + pass + + return "cpu" + + +def poetry_install_with_accel(poetry_exe, script_dir, venv_env, accelerator): + """Run poetry install selecting dependency groups: cuda(default), npu.""" + cmd = [str(poetry_exe), "install"] + print(f"Running poetry install for accelerator={accelerator} -> {' '.join(cmd)}") + subprocess.run( + cmd, + cwd=str(script_dir), + env=venv_env, + check=True, + capture_output=True, + text=True, + ) + + def build(): """ Execute the complete build process. diff --git a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py index 41bc6ec91c8..45db66e0186 100644 --- a/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py +++ b/iotdb-core/ainode/iotdb/ainode/core/manager/utils.py @@ -40,30 +40,6 @@ INFERENCE_EXTRA_MEMORY_RATIO = ( ) # the overhead ratio for inference, used to estimate the pool size -def measure_model_memory(device: torch.device, model_id: str) -> int: - # TODO: support CPU in the future - # TODO: we can estimate the memory usage by running a dummy inference - torch.cuda.empty_cache() - torch.cuda.synchronize(device) - start = torch.cuda.memory_reserved(device) - - model_info = ModelManager().get_model_info(model_id) - model = load_model(model_info).to(device) - torch.cuda.synchronize(device) - end = torch.cuda.memory_reserved(device) - usage = end - start - - # delete model to free memory - del model - torch.cuda.empty_cache() - gc.collect() - - # add inference factor and cuda context overhead - overhead = 500 * 1024**2 # 500 MiB - final = int(max(usage, 1) * INFERENCE_EXTRA_MEMORY_RATIO + overhead) - return final - - def evaluate_system_resources(device: torch.device) -> dict: if device.type == "cuda": free_mem, total_mem = torch.cuda.mem_get_info()
