mirror of
https://github.com/deepmodeling/Uni-Lab-OS
synced 2026-03-24 20:39:20 +00:00
368 lines
13 KiB
Python
368 lines
13 KiB
Python
"""
|
||
环境检查模块
|
||
用于检查并自动安装 UniLabOS 运行所需的 Python 包
|
||
"""
|
||
|
||
import argparse
|
||
import importlib
|
||
import locale
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import List, Optional
|
||
|
||
from unilabos.utils.banner_print import print_status
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 底层安装工具
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _is_chinese_locale() -> bool:
|
||
try:
|
||
lang = locale.getdefaultlocale()[0]
|
||
return bool(lang and ("zh" in lang.lower() or "chinese" in lang.lower()))
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
_USE_UV: Optional[bool] = None
|
||
|
||
|
||
def _has_uv() -> bool:
|
||
global _USE_UV
|
||
if _USE_UV is None:
|
||
_USE_UV = shutil.which("uv") is not None
|
||
return _USE_UV
|
||
|
||
|
||
def _install_packages(
|
||
packages: List[str],
|
||
upgrade: bool = False,
|
||
label: str = "",
|
||
) -> bool:
|
||
"""
|
||
安装/升级一组包。优先 uv pip install，回退 sys pip。
|
||
逐个安装，任意一个失败不影响后续包。
|
||
|
||
Returns:
|
||
True if all succeeded, False otherwise.
|
||
"""
|
||
if not packages:
|
||
return True
|
||
|
||
is_chinese = _is_chinese_locale()
|
||
use_uv = _has_uv()
|
||
failed: List[str] = []
|
||
|
||
for pkg in packages:
|
||
action_word = "升级" if upgrade else "安装"
|
||
if label:
|
||
print_status(f"[{label}] 正在{action_word} {pkg}...", "info")
|
||
else:
|
||
print_status(f"正在{action_word} {pkg}...", "info")
|
||
|
||
if use_uv:
|
||
cmd = ["uv", "pip", "install"]
|
||
if upgrade:
|
||
cmd.append("--upgrade")
|
||
cmd.append(pkg)
|
||
if is_chinese:
|
||
cmd.extend(["--index-url", "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"])
|
||
else:
|
||
cmd = [sys.executable, "-m", "pip", "install"]
|
||
if upgrade:
|
||
cmd.append("--upgrade")
|
||
cmd.append(pkg)
|
||
if is_chinese:
|
||
cmd.extend(["-i", "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"])
|
||
|
||
try:
|
||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
|
||
if result.returncode == 0:
|
||
installer = "uv" if use_uv else "pip"
|
||
print_status(f"✓ {pkg} {action_word}成功 (via {installer})", "success")
|
||
else:
|
||
stderr_short = result.stderr.strip().split("\n")[-1] if result.stderr else "unknown error"
|
||
print_status(f"× {pkg} {action_word}失败: {stderr_short}", "error")
|
||
failed.append(pkg)
|
||
except subprocess.TimeoutExpired:
|
||
print_status(f"× {pkg} {action_word}超时 (300s)", "error")
|
||
failed.append(pkg)
|
||
except Exception as e:
|
||
print_status(f"× {pkg} {action_word}异常: {e}", "error")
|
||
failed.append(pkg)
|
||
|
||
if failed:
|
||
print_status(f"有 {len(failed)} 个包操作失败: {', '.join(failed)}", "error")
|
||
return False
|
||
return True
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# requirements.txt 安装（可多次调用）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def install_requirements_txt(req_path: str | Path, label: str = "") -> bool:
|
||
"""
|
||
读取一个 requirements.txt 文件，检查缺失的包并安装。
|
||
|
||
Args:
|
||
req_path: requirements.txt 文件路径
|
||
label: 日志前缀标签（如 "device_package_sim"）
|
||
|
||
Returns:
|
||
True if all ok, False if any install failed.
|
||
"""
|
||
req_path = Path(req_path)
|
||
if not req_path.exists():
|
||
return True
|
||
|
||
tag = label or req_path.parent.name
|
||
print_status(f"[{tag}] 检查依赖: {req_path}", "info")
|
||
|
||
reqs: List[str] = []
|
||
with open(req_path, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line and not line.startswith("#") and not line.startswith("-"):
|
||
reqs.append(line)
|
||
|
||
if not reqs:
|
||
return True
|
||
|
||
missing: List[str] = []
|
||
for req in reqs:
|
||
pkg_import = req.split(">=")[0].split("==")[0].split("<")[0].split("[")[0].split(">")[0].strip()
|
||
pkg_import = pkg_import.replace("-", "_")
|
||
try:
|
||
importlib.import_module(pkg_import)
|
||
except ImportError:
|
||
missing.append(req)
|
||
|
||
if not missing:
|
||
print_status(f"[{tag}] ✓ 依赖检查通过 ({len(reqs)} 个包)", "success")
|
||
return True
|
||
|
||
print_status(f"[{tag}] 缺失 {len(missing)} 个依赖: {', '.join(missing)}", "warning")
|
||
return _install_packages(missing, label=tag)
|
||
|
||
|
||
def check_device_package_requirements(devices_dirs: list[str]) -> bool:
|
||
"""
|
||
检查 --devices 指定的所有外部设备包目录中的 requirements.txt。
|
||
对每个目录查找 requirements.txt（先在目录内找，再在父目录找）。
|
||
"""
|
||
if not devices_dirs:
|
||
return True
|
||
|
||
all_ok = True
|
||
for d in devices_dirs:
|
||
d_path = Path(d).resolve()
|
||
req_file = d_path / "requirements.txt"
|
||
if not req_file.exists():
|
||
req_file = d_path.parent / "requirements.txt"
|
||
if not req_file.exists():
|
||
continue
|
||
if not install_requirements_txt(req_file, label=d_path.name):
|
||
all_ok = False
|
||
|
||
return all_ok
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# UniLabOS 核心环境检查
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class EnvironmentChecker:
|
||
"""环境检查器"""
|
||
|
||
def __init__(self):
|
||
self.required_packages = {
|
||
"websockets": "websockets",
|
||
"msgcenterpy": "msgcenterpy",
|
||
"orjson": "orjson",
|
||
"opentrons_shared_data": "opentrons_shared_data",
|
||
"typing_extensions": "typing_extensions",
|
||
"crcmod": "crcmod-plus",
|
||
}
|
||
|
||
self.special_packages = {"pylabrobot": "git+https://github.com/Xuwznln/pylabrobot.git"}
|
||
|
||
self.version_requirements = {
|
||
"msgcenterpy": "0.1.8",
|
||
}
|
||
|
||
self.missing_packages: List[tuple] = []
|
||
self.failed_installs: List[tuple] = []
|
||
self.packages_need_upgrade: List[tuple] = []
|
||
|
||
def check_package_installed(self, package_name: str) -> bool:
|
||
try:
|
||
importlib.import_module(package_name)
|
||
return True
|
||
except ImportError:
|
||
return False
|
||
|
||
def get_package_version(self, package_name: str) -> str | None:
|
||
try:
|
||
module = importlib.import_module(package_name)
|
||
return getattr(module, "__version__", None)
|
||
except (ImportError, AttributeError):
|
||
return None
|
||
|
||
def compare_version(self, current: str, required: str) -> bool:
|
||
try:
|
||
current_parts = [int(x) for x in current.split(".")]
|
||
required_parts = [int(x) for x in required.split(".")]
|
||
max_len = max(len(current_parts), len(required_parts))
|
||
current_parts.extend([0] * (max_len - len(current_parts)))
|
||
required_parts.extend([0] * (max_len - len(required_parts)))
|
||
return current_parts >= required_parts
|
||
except Exception:
|
||
return True
|
||
|
||
def check_all_packages(self) -> bool:
|
||
print_status("开始检查环境依赖...", "info")
|
||
|
||
for import_name, pip_name in self.required_packages.items():
|
||
if not self.check_package_installed(import_name):
|
||
self.missing_packages.append((import_name, pip_name))
|
||
elif import_name in self.version_requirements:
|
||
current_version = self.get_package_version(import_name)
|
||
required_version = self.version_requirements[import_name]
|
||
if current_version and not self.compare_version(current_version, required_version):
|
||
print_status(
|
||
f"{import_name} 版本过低 (当前: {current_version}, 需要: >={required_version})",
|
||
"warning",
|
||
)
|
||
self.packages_need_upgrade.append((import_name, pip_name))
|
||
|
||
for package_name, install_url in self.special_packages.items():
|
||
if not self.check_package_installed(package_name):
|
||
self.missing_packages.append((package_name, install_url))
|
||
|
||
all_ok = not self.missing_packages and not self.packages_need_upgrade
|
||
|
||
if all_ok:
|
||
print_status("✓ 所有依赖包检查完成，环境正常", "success")
|
||
return True
|
||
|
||
if self.missing_packages:
|
||
print_status(f"发现 {len(self.missing_packages)} 个缺失的包", "warning")
|
||
if self.packages_need_upgrade:
|
||
print_status(f"发现 {len(self.packages_need_upgrade)} 个需要升级的包", "warning")
|
||
|
||
return False
|
||
|
||
def install_missing_packages(self, auto_install: bool = True) -> bool:
|
||
if not self.missing_packages and not self.packages_need_upgrade:
|
||
return True
|
||
|
||
if not auto_install:
|
||
if self.missing_packages:
|
||
print_status("缺失以下包:", "warning")
|
||
for import_name, pip_name in self.missing_packages:
|
||
print_status(f" - {import_name} ({pip_name})", "warning")
|
||
if self.packages_need_upgrade:
|
||
print_status("需要升级以下包:", "warning")
|
||
for import_name, pip_name in self.packages_need_upgrade:
|
||
print_status(f" - {import_name} ({pip_name})", "warning")
|
||
return False
|
||
|
||
if self.missing_packages:
|
||
pkgs = [pip_name for _, pip_name in self.missing_packages]
|
||
if not _install_packages(pkgs, label="unilabos"):
|
||
self.failed_installs.extend(self.missing_packages)
|
||
|
||
if self.packages_need_upgrade:
|
||
pkgs = [pip_name for _, pip_name in self.packages_need_upgrade]
|
||
if not _install_packages(pkgs, upgrade=True, label="unilabos"):
|
||
self.failed_installs.extend(self.packages_need_upgrade)
|
||
|
||
return not self.failed_installs
|
||
|
||
def verify_installation(self) -> bool:
|
||
if not self.missing_packages and not self.packages_need_upgrade:
|
||
return True
|
||
|
||
print_status("验证安装结果...", "info")
|
||
failed_verification = []
|
||
|
||
for import_name, pip_name in self.missing_packages:
|
||
if not self.check_package_installed(import_name):
|
||
failed_verification.append((import_name, pip_name))
|
||
|
||
for import_name, pip_name in self.packages_need_upgrade:
|
||
if not self.check_package_installed(import_name):
|
||
failed_verification.append((import_name, pip_name))
|
||
elif import_name in self.version_requirements:
|
||
current_version = self.get_package_version(import_name)
|
||
required_version = self.version_requirements[import_name]
|
||
if current_version and not self.compare_version(current_version, required_version):
|
||
failed_verification.append((import_name, pip_name))
|
||
print_status(
|
||
f" {import_name} 版本仍然过低 (当前: {current_version}, 需要: >={required_version})",
|
||
"error",
|
||
)
|
||
|
||
if failed_verification:
|
||
print_status(f"有 {len(failed_verification)} 个包验证失败:", "error")
|
||
for import_name, pip_name in failed_verification:
|
||
print_status(f" - {import_name}", "error")
|
||
return False
|
||
|
||
print_status("✓ 所有包验证通过", "success")
|
||
return True
|
||
|
||
|
||
def check_environment(auto_install: bool = True, show_details: bool = True) -> bool:
|
||
"""
|
||
检查环境并自动安装缺失的包
|
||
|
||
Args:
|
||
auto_install: 是否自动安装缺失的包
|
||
show_details: 是否显示详细信息
|
||
|
||
Returns:
|
||
bool: 环境检查是否通过
|
||
"""
|
||
checker = EnvironmentChecker()
|
||
|
||
if checker.check_all_packages():
|
||
return True
|
||
|
||
if not checker.install_missing_packages(auto_install):
|
||
if show_details:
|
||
print_status("请手动安装缺失的包后重新启动程序", "error")
|
||
return False
|
||
|
||
if not checker.verify_installation():
|
||
if show_details:
|
||
print_status("安装验证失败，请检查网络连接或手动安装", "error")
|
||
return False
|
||
|
||
return True
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="UniLabOS 环境依赖检查工具")
|
||
parser.add_argument("--no-auto-install", action="store_true", help="仅检查环境，不自动安装缺失的包")
|
||
parser.add_argument("--silent", action="store_true", help="静默模式，不显示详细信息")
|
||
|
||
args = parser.parse_args()
|
||
|
||
auto_install = not args.no_auto_install
|
||
show_details = not args.silent
|
||
|
||
success = check_environment(auto_install=auto_install, show_details=show_details)
|
||
|
||
if not success:
|
||
if show_details:
|
||
print_status("环境检查失败", "error")
|
||
sys.exit(1)
|
||
else:
|
||
if show_details:
|
||
print_status("环境检查完成", "success")
|