forked from ci4s/aim-proxy
58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
from aim.sdk.repo import Repo
|
|
import httpx
|
|
from datetime import datetime
|
|
import uvicorn
|
|
import traceback
|
|
|
|
repo_path = 'aim://172.20.32.181:30038'
|
|
#repo_path = 'aim://my-aim-server-service.aim:53800'
|
|
# host = 'http://my-aim-ui-service.aim:43800'
|
|
go_service_host = 'http://172.20.32.186:7124' # Go服务的地址
|
|
|
|
# 打开 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)
|
|
if paginated and limit:
|
|
runs = runs[:limit]
|
|
|
|
indexed_runs = [{"index": i, "hash": run.hash, "creation_time": run.creation_time} for i, run in enumerate(runs)]
|
|
print("indexed_runs:", indexed_runs)
|
|
# 将indexed_runs传递给Go服务以获取详细信息
|
|
http_start_time = datetime.now().timestamp()
|
|
async with httpx.AsyncClient(timeout=20) as client:
|
|
response = await client.post(f"{go_service_host}/api/fetch_run_info", json=indexed_runs)
|
|
response.raise_for_status()
|
|
results = response.json()
|
|
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)
|