forked from ci4s/aim-proxy
92 lines
3.4 KiB
Python
92 lines
3.4 KiB
Python
from fastapi import FastAPI, HTTPException
|
||
from aim.sdk.repo import Repo
|
||
import httpx
|
||
import traceback
|
||
import uvicorn
|
||
from datetime import datetime
|
||
import asyncio
|
||
|
||
# AIM repository路径和目标服务的主机地址
|
||
repo_path = 'aim://172.20.32.181:30058'
|
||
host = 'http://172.20.32.181:30059'
|
||
|
||
# 打开 AIM repository
|
||
repo = Repo(path=repo_path)
|
||
|
||
# 创建 FastAPI 应用实例
|
||
app = FastAPI()
|
||
|
||
# 定义异步函数，用于获取每个 run 的详细信息
|
||
async def fetch_run_info(run, results):
|
||
try:
|
||
# 构建获取信息的 URL
|
||
url = f"{host}/api/runs/{run.hash}/info"
|
||
# 使用 httpx.AsyncClient 发起异步 GET 请求
|
||
async with httpx.AsyncClient() as client:
|
||
response = await client.get(url)
|
||
response.raise_for_status()
|
||
info = response.json()
|
||
info['run_hash'] = run.hash # 将 run_hash 添加到返回的信息中
|
||
results.append(info) # 将信息添加到结果列表中
|
||
except Exception as e:
|
||
traceback.print_exc() # 打印异常堆栈信息
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
# 定义路由 /api/runs/search/run，接收查询参数并发起调用
|
||
@app.get("/api/runs/search/run")
|
||
async def search_runs(query: str = None, paginated: bool = True, offset: str = None, limit: int = None):
|
||
try:
|
||
# 记录查询开始时间
|
||
query_start_time = datetime.now().timestamp()
|
||
# 使用 repo.query_runs 查询符合条件的 runs，并获取迭代器
|
||
query_res = repo.query_runs(query=query, paginated=paginated, offset=offset).iter_runs()
|
||
# 将查询结果中的 run 对象提取到列表中
|
||
runs = [item.run for item in query_res]
|
||
# 记录查询结束时间和持续时间
|
||
query_end_time = datetime.now().timestamp()
|
||
query_duration = query_end_time - query_start_time
|
||
print("query time is ", query_duration)
|
||
|
||
# 根据 run.creation_time 降序排序 runs 列表
|
||
runs = sorted(runs, key=lambda x: x.creation_time, reverse=True)
|
||
|
||
results = [] # 初始化结果列表
|
||
|
||
# 记录 HTTP 请求开始时间
|
||
http_start_time = datetime.now().timestamp()
|
||
|
||
tasks = [] # 初始化任务列表
|
||
for run in runs:
|
||
# 创建异步任务，并添加到任务列表
|
||
task = fetch_run_info(run, results)
|
||
tasks.append(task)
|
||
# 当任务数达到 10 时，使用 asyncio.gather 并发执行这些任务
|
||
if len(tasks) >= 8:
|
||
await asyncio.gather(*tasks)
|
||
tasks = [] # 清空任务列表，以便下一批任务
|
||
|
||
# 等待剩余的任务完成
|
||
if tasks:
|
||
await asyncio.gather(*tasks)
|
||
|
||
# 记录 HTTP 请求结束时间和持续时间
|
||
http_end_time = datetime.now().timestamp()
|
||
http_duration = http_end_time - http_start_time
|
||
print("http time is ", http_duration)
|
||
|
||
# 如果是分页模式，并且有限制结果数量，则截取结果列表
|
||
if paginated:
|
||
results = results[:limit] if limit else results
|
||
|
||
# 返回最终结果
|
||
print("results len is ", len(results))
|
||
return results
|
||
|
||
except Exception as e:
|
||
traceback.print_exc() # 打印异常堆栈信息
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
# 如果是主程序入口，则运行 FastAPI 应用
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=7123)
|