200 lines
6.5 KiB
Python
200 lines
6.5 KiB
Python
from aim.storage.treeutils import decode_tree
|
||
|
||
from fastapi import FastAPI, HTTPException, Request
|
||
import httpx
|
||
import traceback
|
||
import uvicorn
|
||
|
||
import json
|
||
from typing import Iterator
|
||
import struct
|
||
|
||
from pydantic import BaseModel
|
||
from typing import List
|
||
from aim.sdk.repo import Repo
|
||
|
||
import os
|
||
import logging
|
||
import time
|
||
|
||
app = FastAPI(
|
||
title="aim-proxy API",
|
||
description="This is aim-proxy API",
|
||
version="1.0.0",
|
||
openapi_url="/api/v1/openapi.json",
|
||
docs_url="/api/v1/docs",
|
||
redoc_url="/api/v1/redoc"
|
||
)
|
||
|
||
# 设置日志格式
|
||
logging.basicConfig(level=logging.INFO)
|
||
|
||
|
||
@app.middleware("http")
|
||
async def log_request_time(request: Request, call_next):
|
||
start_time = time.time() # 记录请求开始时间
|
||
response = await call_next(request) # 处理请求
|
||
process_time = time.time() - start_time # 计算处理时间
|
||
logging.info(f"Request {request.method} {request.url} completed in {process_time:.4f} seconds")
|
||
return response # 返回响应
|
||
|
||
|
||
def is_running_in_kubernetes():
|
||
return 'KUBERNETES_SERVICE_HOST' in os.environ and 'KUBERNETES_SERVICE_PORT' in os.environ
|
||
|
||
|
||
if is_running_in_kubernetes():
|
||
print("Running in Kubernetes environment")
|
||
repo_path = 'aim://aim-server-service.aim:53800'
|
||
host = 'http://aim-ui-service.aim:43800'
|
||
else:
|
||
repo_path = 'aim://172.20.32.197:30058'
|
||
host = 'http://172.20.32.197:30059'
|
||
print("Not running in Kubernetes environment")
|
||
|
||
api_url = f'{host}/api/runs/search/run/'
|
||
delete_url = f'{host}/api/runs/delete-batch'
|
||
# 打开 Aim repository
|
||
repo = Repo(path=repo_path)
|
||
|
||
|
||
@app.get("/api/runs/search/run")
|
||
async def search_runs(query: str = None, offset: str = None, limit: int = 45):
|
||
try:
|
||
print("query is ", query)
|
||
print("offset is ", offset)
|
||
print("limit is ", limit)
|
||
#query_url = f'{api_url}?q={query}&limit={limit}&offset={offset}'
|
||
print("query_url is ", api_url)
|
||
data = await fetch_run_data(api_url, query=query, limit=limit, offset=offset)
|
||
result = parse_run_data(data)
|
||
return result
|
||
|
||
except Exception as e:
|
||
print(f"An error occurred while searching runs: {str(e)}")
|
||
traceback.print_exc()
|
||
raise HTTPException(status_code=500, detail=str(e))
|
||
|
||
|
||
# 定义请求体模型
|
||
class RunDeleteRequest(BaseModel):
|
||
run_hashes: List[str]
|
||
|
||
|
||
@app.post("/api/runs/delete-batch/run")
|
||
async def delete_runs(request: RunDeleteRequest):
|
||
try:
|
||
if not request.run_hashes:
|
||
raise HTTPException(status_code=400, detail="run_hash is required")
|
||
|
||
async with httpx.AsyncClient() as client:
|
||
run_hashes = request.run_hashes
|
||
print("delete url is ", delete_url)
|
||
print("run_hashes is ", run_hashes)
|
||
response = await client.post(delete_url, json=run_hashes)
|
||
response.raise_for_status() # 如果响应状态码不是 200，则引发异常
|
||
return {"message": "Runs deleted successfully"}
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
print(f"HTTP error occurred: {e.response.status_code}")
|
||
print(e.response.text)
|
||
except Exception as e:
|
||
print(f"An error occurred: {str(e)}")
|
||
|
||
|
||
def decode_encoded_tree_stream(stream: Iterator[bytes], concat_chunks=False) -> bytes:
|
||
# TODO: handle case when chunk ends at the middle of key/value
|
||
# TODO: if remaining part of chunk cannot be unpacked, prepend to next one and try with new chunk
|
||
prev_chunk_tail = b''
|
||
if concat_chunks:
|
||
data = b''
|
||
for chunk in stream:
|
||
data += chunk
|
||
while data:
|
||
(key_size,), data_tail = struct.unpack('I', data[:4]), data[4:]
|
||
key, data_tail = data_tail[:key_size], data_tail[key_size:]
|
||
|
||
(value_size,), data_tail = struct.unpack('I', data_tail[:4]), data_tail[4:]
|
||
value, data_tail = data_tail[:value_size], data_tail[value_size:]
|
||
data = data_tail
|
||
yield key, value
|
||
else:
|
||
for chunk in stream:
|
||
data = prev_chunk_tail + chunk
|
||
prev_chunk_tail = b''
|
||
while data:
|
||
try:
|
||
(key_size,), data_tail = struct.unpack('I', data[:4]), data[4:]
|
||
key, data_tail = data_tail[:key_size], data_tail[key_size:]
|
||
|
||
(value_size,), data_tail = struct.unpack('I', data_tail[:4]), data_tail[4:]
|
||
value, data_tail = data_tail[:value_size], data_tail[value_size:]
|
||
data = data_tail
|
||
except Exception:
|
||
prev_chunk_tail = data
|
||
break
|
||
|
||
yield key, value
|
||
|
||
assert prev_chunk_tail == b''
|
||
|
||
|
||
async def fetch_run_data(api_url: str, query: str = '', limit: int = 45, offset: str = None) -> dict:
|
||
try:
|
||
params = {
|
||
'q': query,
|
||
'limit': limit,
|
||
'offset': offset
|
||
}
|
||
async with httpx.AsyncClient() as client:
|
||
response = await client.get(api_url, params=params)
|
||
response.raise_for_status() # 如果响应状态码不是 200，则引发异常
|
||
decoded_response = decode_tree(decode_encoded_tree_stream(response.iter_bytes(chunk_size=512 * 1024)))
|
||
return decoded_response
|
||
except httpx.HTTPStatusError as e:
|
||
print(f"HTTP error occurred: {e.response.status_code}")
|
||
print(e.response.text)
|
||
except Exception as e:
|
||
print(f"An error occurred: {str(e)}")
|
||
|
||
|
||
def process_metrics(hash, data):
|
||
# 过滤掉名字以 __system 开头的 metric
|
||
filtered_metrics = [metric for metric in data['traces']['metric'] if not metric['name'].startswith('__system')]
|
||
|
||
# 给剩余的 metric 添加 last_value 键
|
||
for metric in filtered_metrics:
|
||
metric['last_value'] = metric['values']['last']
|
||
|
||
# 更新数据中的 metric 列表
|
||
data['traces']['metric'] = filtered_metrics
|
||
data['run_hash'] = hash
|
||
return data
|
||
|
||
|
||
def parse_run_data(decoded_response):
|
||
try:
|
||
if not decoded_response:
|
||
print("No data to parse")
|
||
return
|
||
result = []
|
||
for i, run in decoded_response.items():
|
||
print(f"Run {i}:")
|
||
if i.startswith('progress'):
|
||
continue
|
||
|
||
run = process_metrics(i, run)
|
||
print("data is ", json.dumps(run, indent=4))
|
||
result.append(run)
|
||
|
||
print("result is ", json.dumps(result, indent=4))
|
||
return result
|
||
|
||
|
||
except Exception as e:
|
||
print(f"An error occurred while parsing data: {str(e)}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=7123)
|