forked from ci4s/aim-proxy
使用更高效的api, 节约用时
This commit is contained in:
parent
16c1fb7b7b
commit
32be98f230
|
|
@ -12,6 +12,6 @@ RUN pip install httpx -i https://pypi.tuna.tsinghua.edu.cn/simple
|
|||
ENV http_proxy=""
|
||||
ENV https_proxy=""
|
||||
|
||||
COPY server.py /app/server.py
|
||||
COPY aim_proxy.py /app/server.py
|
||||
|
||||
ENTRYPOINT ["python", "server.py"]
|
||||
|
|
@ -0,0 +1,192 @@
|
|||
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()
|
||||
|
||||
# 设置日志格式
|
||||
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://my-aim-server-service.aim:53800'
|
||||
host = 'http://my-aim-ui-service.aim:43800'
|
||||
else:
|
||||
repo_path = 'aim://172.20.32.181:30058'
|
||||
host = 'http://172.20.32.181: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 ", query_url)
|
||||
data = await fetch_run_data(query_url)
|
||||
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):
|
||||
try:
|
||||
params = {
|
||||
'q': query,
|
||||
'limit': limit
|
||||
}
|
||||
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)
|
||||
Loading…
Reference in New Issue