112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from .config import infer_kind, load_config
|
|
|
|
|
|
PROJECT_DIR = Path(__file__).resolve().parent.parent
|
|
DEFAULT_CONFIG = PROJECT_DIR / "config" / "config.json"
|
|
|
|
|
|
def _parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(prog="localpilot", description="LocalPilot 本地模型后端")
|
|
parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="配置 JSON 路径")
|
|
sub = parser.add_subparsers(dest="command", required=True)
|
|
sub.add_parser("models", help="列出配置模型")
|
|
sub.add_parser("doctor", help="检查本机运行环境")
|
|
serve = sub.add_parser("serve", help="启动 OpenAI 兼容 API")
|
|
serve.add_argument("--host")
|
|
serve.add_argument("--port", type=int)
|
|
sub.add_parser("tui", help="启动 Textual TUI")
|
|
load = sub.add_parser("load", help="预加载模型")
|
|
load.add_argument("model_id")
|
|
inspect = sub.add_parser("inspect", help="推断模型格式")
|
|
inspect.add_argument("path")
|
|
return parser
|
|
|
|
|
|
def _doctor() -> dict[str, object]:
|
|
result: dict[str, object] = {"python": sys.version, "executable": sys.executable}
|
|
try:
|
|
import torch
|
|
|
|
result["torch"] = torch.__version__
|
|
result["cuda_available"] = bool(torch.cuda.is_available())
|
|
result["cuda_device"] = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
|
|
result["cuda_version"] = torch.version.cuda
|
|
except Exception as exc:
|
|
result["torch_error"] = str(exc)
|
|
try:
|
|
import onnxruntime as ort
|
|
|
|
result["onnxruntime"] = ort.__version__
|
|
result["onnx_providers"] = ort.get_available_providers()
|
|
except Exception as exc:
|
|
result["onnx_error"] = str(exc)
|
|
result["conda_prefix"] = os.getenv("CONDA_PREFIX")
|
|
return result
|
|
|
|
|
|
def main() -> None:
|
|
args = _parser().parse_args()
|
|
config_path = Path(args.config)
|
|
config = load_config(config_path)
|
|
|
|
if args.command == "models":
|
|
print(json.dumps([
|
|
{"id": model.id, "kind": infer_kind(model), "path": model.path, "enabled": model.enabled}
|
|
for model in config.models
|
|
if model.enabled
|
|
], ensure_ascii=False, indent=2))
|
|
return
|
|
if args.command == "doctor":
|
|
print(json.dumps(_doctor(), ensure_ascii=False, indent=2))
|
|
return
|
|
if args.command == "inspect":
|
|
from .config import ModelConfig
|
|
|
|
spec = ModelConfig(id=Path(args.path).stem, path=args.path)
|
|
print(json.dumps({"path": args.path, "kind": infer_kind(spec)}, ensure_ascii=False, indent=2))
|
|
return
|
|
if args.command == "serve":
|
|
import uvicorn
|
|
|
|
from .server import create_app
|
|
|
|
app = create_app(config_path)
|
|
uvicorn.run(
|
|
app,
|
|
host=args.host or config.runtime.host,
|
|
port=args.port or config.runtime.port,
|
|
log_level="info",
|
|
)
|
|
return
|
|
if args.command == "tui":
|
|
from .tui import LocalPilotTUI
|
|
|
|
LocalPilotTUI(config).run()
|
|
return
|
|
if args.command == "load":
|
|
from .manager import ModelManager
|
|
|
|
async def run() -> None:
|
|
manager = ModelManager(config)
|
|
try:
|
|
backend = await manager.load(args.model_id)
|
|
print(json.dumps(backend.info(), ensure_ascii=False, indent=2))
|
|
finally:
|
|
await manager.shutdown()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|