forked from ci4s/aim-proxy
90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
from fastapi import FastAPI, HTTPException
|
|
from aim.sdk.repo import Repo
|
|
import httpx
|
|
import traceback
|
|
from datetime import datetime
|
|
import uvicorn
|
|
|
|
repo_path = 'aim://my-aim-server-service.aim:53800'
|
|
#repo_path = 'aim://172.20.32.181:30058'
|
|
host = 'http://my-aim-ui-service.aim:43800'
|
|
#host = 'http://172.20.32.181:30059'
|
|
repo = Repo(path=repo_path)
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
async def fetch_run_info(client, host, run_hash):
|
|
try:
|
|
url = f"{host}/api/runs/{run_hash}/info"
|
|
response = await client.get(url)
|
|
response.raise_for_status()
|
|
info = response.json()
|
|
return info
|
|
except httpx.HTTPStatusError as e:
|
|
raise HTTPException(status_code=e.response.status_code, detail=str(e))
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
async def fetch_run_metrics(client, host, run_hash, data):
|
|
try:
|
|
url = f"{host}/api/runs/{run_hash}/metric/get-batch"
|
|
response = await client.post(url, json=data)
|
|
response.raise_for_status()
|
|
metrics = response.json()
|
|
return metrics
|
|
except httpx.HTTPStatusError as e:
|
|
raise HTTPException(status_code=e.response.status_code, detail=str(e))
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
def extract_latest_metrics(metrics, name, context):
|
|
for metric in metrics:
|
|
if metric["name"] == name and metric["context"] == context:
|
|
return metric["values"][-1] if metric["values"] else None
|
|
return None
|
|
|
|
|
|
@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()
|
|
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)
|
|
|
|
runs = sorted(runs, key=lambda x: x.creation_time, reverse=True)
|
|
if paginated and limit:
|
|
runs = runs[:limit]
|
|
|
|
results = []
|
|
async with httpx.AsyncClient() as client:
|
|
for run in runs:
|
|
run_info = await fetch_run_info(client, host, run.hash)
|
|
data = run_info["traces"]["metric"]
|
|
run_metrics = await fetch_run_metrics(client, host, run.hash, data)
|
|
for metric in run_metrics:
|
|
metric["last_value"] = metric["values"][-1] if metric["values"] else None
|
|
|
|
run_info["traces"]["metric"] = run_metrics
|
|
run_info["run_hash"] = run.hash
|
|
results.append(run_info)
|
|
|
|
http_end_time = datetime.now().timestamp()
|
|
http_duration = http_end_time - query_end_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)
|