forked from ci4s/aim-proxy
56 lines
2.2 KiB
Python
56 lines
2.2 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
from aim.sdk.repo import Repo
|
|
import httpx
|
|
from aim.sdk.types import QueryReportMode
|
|
import traceback
|
|
import uvicorn
|
|
from datetime import datetime
|
|
#repo_path = 'aim://172.20.32.181:30058'
|
|
repo_path = 'aim://my-aim-server-service.aim:53800'
|
|
host = 'http://my-aim-ui-service.aim:43800'
|
|
# 打开 Aim repository
|
|
repo = Repo(path=repo_path)
|
|
|
|
app = FastAPI()
|
|
@app.get("/api/runs/search/run")
|
|
async def search_runs(query: str = None, paginated: bool = True, offset: str = None, limit: int = None):
|
|
try:
|
|
print("query is ", query)
|
|
print("paginated is ", paginated)
|
|
print("offset is ", offset)
|
|
print("limit is ", limit)
|
|
query_start_time = datetime.now().timestamp()
|
|
query_res = repo.query_runs(query=query, paginated=paginated, offset=offset).iter_runs()
|
|
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 = sorted(runs, key=lambda x: x.creation_time, reverse=True)
|
|
results = []
|
|
http_start_time = datetime.now().timestamp()
|
|
for run in runs:
|
|
# 调用获取详细信息的接口
|
|
url = f"{host}/api/runs/{run.hash}/info"
|
|
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)
|
|
if paginated:
|
|
if limit and len(results) >= limit:
|
|
break
|
|
print("result len is ", len(results))
|
|
http_end_time = datetime.now().timestamp()
|
|
http_duration = http_end_time - http_start_time
|
|
print("http time is ", http_duration)
|
|
return results
|
|
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=7123)
|