aim-proxy/server_new.py

107 lines
3.2 KiB
Python

from fastapi import FastAPI, HTTPException
from aim.sdk.repo import Repo
import httpx
from aim.sdk.types import QueryReportMode
from typing import Optional, Tuple, Dict, Any
from aim.sdk.repo import Repo
from aim.sdk.run import Run
from aim.web.api.runs.utils import (
checked_query,
collect_requested_metric_traces,
convert_nan_and_inf_to_str,
custom_aligned_metrics_streamer,
get_project_repo,
get_run_or_404,
get_run_params,
get_run_props,
get_run_artifacts,
metric_search_result_streamer,
run_active_result_streamer,
run_search_result_streamer,
run_logs_streamer,
run_log_records_streamer,
)
app = FastAPI()
repo_path = 'aim://172.20.32.181:30058'
# 打开 Aim repository
repo = Repo(path=repo_path)
@app.get("/api/runs/search/run")
async def search_runs(query: str, paginated: bool = True, offset: str = None, limit: int = None):
print("query is ", query)
print("paginated is ", paginated)
print("offset is ", offset)
print("limit is ", limit)
query_res = repo.query_runs(query=query, paginated=paginated, offset=offset).iter_runs()
runs = [item.run for item in query_res]
# 根据run.creation_time 降序排序
runs = sorted(runs, key=lambda x: x.creation_time, reverse=True)
results = []
print("runs len is ", len(runs))
for run in runs:
# sequence = repo.available_sequence_types()
# sequence = ["metric"]
# print("sequence is :", sequence)
# response = {
# 'params': get_run_params(run, skip_system=True),
# 'traces': run.collect_sequence_info(sequence, skip_last_value=True),
# # 'props': get_run_props(run),
# 'artifacts': get_run_artifacts(run),
# }
# # Convert NaN and Inf to strings
# response = convert_nan_and_inf_to_str(response)
#
# # response['props'].update({
# # 'notes': len(run.props.notes_obj)
# # })
# results.append(response)
re = get_run_info2(run.hash, repo, skip_system=True, sequence=("metric",))
re['run_hash'] = run.hash
results.append(re)
print("result len is ", len(results))
return results
def get_run_info2(run_id: str,
repo: Repo,
skip_system: Optional[bool] = False,
sequence: Optional[Tuple[str, ...]] = ()) -> Dict[str, Any]:
run = get_run_or_404(run_id, repo=repo)
if sequence != ():
try:
repo.validate_sequence_types(sequence)
except ValueError as e:
raise ValueError(f"Invalid sequence types: {str(e)}")
else:
sequence = repo.available_sequence_types()
response = {
'params': get_run_params(run, skip_system=skip_system),
'traces': run.collect_sequence_info(sequence, skip_last_value=False),
# 'props': get_run_props(run),
'artifacts': get_run_artifacts(run),
}
# Convert NaN and Inf to strings
response = convert_nan_and_inf_to_str(response)
# response['props'].update({
# 'notes': len(run.props.notes_obj)
# })
return response
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7124)