wehub-resource-sync
599 lines
20 KiB
Python
599 lines
20 KiB
Python
import logging
|
|
import os
|
|
import random
|
|
import traceback
|
|
from collections import defaultdict
|
|
from contextlib import redirect_stderr, redirect_stdout
|
|
from datetime import datetime, timezone
|
|
from multiprocessing import Process, Queue
|
|
from pathlib import Path
|
|
from queue import Empty
|
|
|
|
import randomname
|
|
import typer
|
|
from flask import Flask, jsonify, request, send_file, send_from_directory
|
|
from flask_cors import CORS
|
|
from werkzeug.utils import secure_filename
|
|
|
|
from rdagent.log.storage import FileStorage
|
|
from rdagent.log.ui.conf import UI_SETTING
|
|
from rdagent.log.ui.storage import WebStorage
|
|
|
|
app = Flask(__name__, static_folder=str(Path(UI_SETTING.static_path).resolve()))
|
|
CORS(app)
|
|
app.config["UI_SERVER_PORT"] = 19899
|
|
|
|
_YELLOW = "\033[33m"
|
|
_RESET = "\033[0m"
|
|
|
|
|
|
class _YellowWarningFormatter(logging.Formatter):
|
|
def format(self, record: logging.LogRecord) -> str:
|
|
if record.levelno == logging.WARNING:
|
|
record.levelname = f"{_YELLOW}{record.levelname}{_RESET}"
|
|
return super().format(record)
|
|
|
|
|
|
def _configure_app_logger() -> None:
|
|
formatter = _YellowWarningFormatter(
|
|
fmt="[%(asctime)s] %(levelname)s in %(module)s: %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
for handler in app.logger.handlers:
|
|
handler.setFormatter(formatter)
|
|
|
|
|
|
_configure_app_logger()
|
|
|
|
|
|
_TARGETS_WITHOUT_USER_INTERACTION = {"general_model", "fin_factor_report"}
|
|
|
|
|
|
class RDAgentTask:
|
|
def __init__(
|
|
self,
|
|
target_name: str,
|
|
kwargs: dict,
|
|
stdout_path: str,
|
|
log_trace_path: str,
|
|
scenario: str,
|
|
trace_name: str,
|
|
ui_server_port: int | None = None,
|
|
create_process: bool = True,
|
|
) -> None:
|
|
self.target_name = target_name
|
|
self.kwargs = kwargs
|
|
self.stdout_path = stdout_path
|
|
self.log_trace_path = log_trace_path
|
|
self.scenario = scenario
|
|
self.trace_name = trace_name
|
|
self.ui_server_port = ui_server_port
|
|
self.process: Process | None = None
|
|
|
|
# Two IPC queues for user interaction.
|
|
# - `user_request_q`: rdagent subprocess -> server (dicts to render on frontend)
|
|
# - `user_response_q`: server -> rdagent subprocess (user input dicts)
|
|
# NOTE: Use multiprocessing.Queue because rdagent is started as a separate process.
|
|
self.user_request_q: Queue = Queue(maxsize=1024)
|
|
self.user_response_q: Queue = Queue(maxsize=1024)
|
|
|
|
if create_process:
|
|
self.process = Process(
|
|
target=self._run,
|
|
name=f"rdagent:{self.scenario}:{self.trace_name}",
|
|
)
|
|
self.messages: list[dict] = []
|
|
self.pointers: defaultdict[str, int] = defaultdict(int)
|
|
|
|
def start(self) -> None:
|
|
if self.process is not None:
|
|
self.process.start()
|
|
|
|
def is_alive(self) -> bool:
|
|
return self.process is not None and self.process.is_alive()
|
|
|
|
def get_end_code(self) -> int:
|
|
if self.process is None or self.process.exitcode is None:
|
|
return 0
|
|
return self.process.exitcode
|
|
|
|
def stop(self) -> None:
|
|
if self.process is not None and self.process.is_alive():
|
|
self.process.terminate()
|
|
self.process.join()
|
|
|
|
# Best-effort cleanup for IPC queues.
|
|
for q in (self.user_request_q, self.user_response_q):
|
|
try:
|
|
q.cancel_join_thread()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
q.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def _run(self) -> None:
|
|
from rdagent.log.conf import LOG_SETTINGS
|
|
|
|
LOG_SETTINGS.set_ui_server_port(self.ui_server_port)
|
|
|
|
from rdagent.log import rdagent_logger
|
|
|
|
rdagent_logger.refresh_storages_from_settings()
|
|
rdagent_logger.set_storages_path(self.log_trace_path)
|
|
Path(self.stdout_path).parent.mkdir(parents=True, exist_ok=True)
|
|
with open(self.stdout_path, "w") as log_file:
|
|
with redirect_stdout(log_file), redirect_stderr(log_file):
|
|
rdagent_logger.rebind_console_to_current_streams()
|
|
try:
|
|
# Only interactive targets should receive IPC queues.
|
|
if self.target_name not in _TARGETS_WITHOUT_USER_INTERACTION:
|
|
self.kwargs.setdefault(
|
|
"user_interaction_queues",
|
|
(self.user_request_q, self.user_response_q),
|
|
)
|
|
|
|
if self.target_name == "data_science":
|
|
from rdagent.app.data_science.loop import main as data_science
|
|
|
|
data_science(**self.kwargs)
|
|
elif self.target_name == "general_model":
|
|
from rdagent.app.general_model.general_model import (
|
|
extract_models_and_implement as general_model,
|
|
)
|
|
|
|
general_model(**self.kwargs)
|
|
elif self.target_name == "fin_factor":
|
|
from rdagent.app.qlib_rd_loop.factor import main as fin_factor
|
|
|
|
fin_factor(**self.kwargs)
|
|
elif self.target_name == "fin_factor_report":
|
|
from rdagent.app.qlib_rd_loop.factor_from_report import (
|
|
main as fin_factor_report,
|
|
)
|
|
|
|
fin_factor_report(**self.kwargs)
|
|
elif self.target_name == "fin_model":
|
|
from rdagent.app.qlib_rd_loop.model import main as fin_model
|
|
|
|
fin_model(**self.kwargs)
|
|
elif self.target_name == "fin_quant":
|
|
from rdagent.app.qlib_rd_loop.quant import main as fin_quant
|
|
|
|
fin_quant(**self.kwargs)
|
|
else:
|
|
raise ValueError(f"Unknown target: {self.target_name}")
|
|
except Exception:
|
|
traceback.print_exc()
|
|
|
|
|
|
rdagent_processes: dict[str, RDAgentTask] = {}
|
|
log_folder_path = Path(UI_SETTING.trace_folder).absolute()
|
|
|
|
|
|
def _drain_user_requests_into_messages(task: RDAgentTask) -> None:
|
|
"""Move a single pending user-interaction request into `task.messages`.
|
|
|
|
Assumption: each rdagent process only has one active request at a time.
|
|
"""
|
|
|
|
try:
|
|
req = task.user_request_q.get_nowait()
|
|
except Empty:
|
|
return
|
|
except Exception:
|
|
return
|
|
|
|
# Standardize the message shape for the frontend.
|
|
# The agent can send either a full message dict, or a raw content dict.
|
|
if isinstance(req, dict) and {"tag", "timestamp", "content"}.issubset(req.keys()):
|
|
msg = req
|
|
else:
|
|
msg = {
|
|
"tag": "user_interaction.request",
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"content": req,
|
|
}
|
|
task.messages.append(msg)
|
|
|
|
|
|
@app.route("/favicon.ico")
|
|
def favicon():
|
|
return send_from_directory(app.static_folder, "favicon.ico", mimetype="image/vnd.microsoft.icon")
|
|
|
|
|
|
def _normalize_static_request_path(fn: str) -> str:
|
|
static_prefix = UI_SETTING.static_path.strip("./")
|
|
if static_prefix and fn.startswith(f"{static_prefix}/"):
|
|
return fn[len(static_prefix) + 1 :]
|
|
return fn
|
|
|
|
|
|
def _get_or_create_task(trace_id: str) -> RDAgentTask:
|
|
task = rdagent_processes.get(trace_id)
|
|
if task is None:
|
|
task = RDAgentTask(
|
|
target_name="",
|
|
kwargs={},
|
|
stdout_path="",
|
|
log_trace_path=trace_id,
|
|
scenario="",
|
|
trace_name="",
|
|
ui_server_port=None,
|
|
create_process=False,
|
|
)
|
|
rdagent_processes[trace_id] = task
|
|
return task
|
|
|
|
|
|
def _resolve_stdout_path(trace_id: str) -> Path | None:
|
|
normalized_trace_id = str(trace_id or "").strip()
|
|
if not normalized_trace_id:
|
|
return None
|
|
|
|
task = rdagent_processes.get(str(log_folder_path / normalized_trace_id))
|
|
if task is None or not task.stdout_path:
|
|
return None
|
|
|
|
stdout_path = Path(task.stdout_path).resolve()
|
|
|
|
try:
|
|
if os.path.commonpath([str(stdout_path), str(log_folder_path)]) != str(log_folder_path):
|
|
return None
|
|
except ValueError:
|
|
return None
|
|
|
|
return stdout_path
|
|
|
|
|
|
def read_trace(log_path: Path, id: str = "") -> None:
|
|
fs = FileStorage(log_path)
|
|
ws = WebStorage(port=1, path=log_path)
|
|
task = _get_or_create_task(id)
|
|
task.messages = []
|
|
last_timestamp = None
|
|
for msg in fs.iter_msg():
|
|
data = ws._obj_to_json(obj=msg.content, tag=msg.tag, id=id, timestamp=msg.timestamp.isoformat())
|
|
if data:
|
|
if isinstance(data, list):
|
|
for d in data:
|
|
task.messages.append(d["msg"])
|
|
last_timestamp = msg.timestamp
|
|
else:
|
|
task.messages.append(data["msg"])
|
|
last_timestamp = msg.timestamp
|
|
|
|
now = datetime.now(timezone.utc)
|
|
if last_timestamp and (now - last_timestamp).total_seconds() > 1800:
|
|
task.messages.append(
|
|
{
|
|
"tag": "END",
|
|
"timestamp": now.isoformat(),
|
|
"content": {"error_msg": "Trace session has ended.", "end_code": 0},
|
|
}
|
|
)
|
|
|
|
|
|
def _collect_existing_trace_ids(trace_root: Path) -> list[str]:
|
|
"""Return trace ids that should be visible in the UI history panel."""
|
|
|
|
if not trace_root.exists():
|
|
return []
|
|
|
|
trace_ids: list[str] = []
|
|
for trace_dir in sorted(trace_root.glob("*/*"), key=lambda p: str(p)):
|
|
if not trace_dir.is_dir():
|
|
continue
|
|
if "uploads" in trace_dir.relative_to(trace_root).parts:
|
|
continue
|
|
if not any(trace_dir.rglob("*.pkl")):
|
|
continue
|
|
|
|
trace_ids.append(trace_dir.relative_to(trace_root).as_posix())
|
|
|
|
return trace_ids
|
|
|
|
|
|
def _load_existing_traces(trace_root: Path) -> None:
|
|
"""Load persisted traces into memory so the UI survives a server restart."""
|
|
|
|
for trace_id in _collect_existing_trace_ids(trace_root):
|
|
trace_dir = trace_root / trace_id
|
|
|
|
try:
|
|
read_trace(trace_dir, id=str(trace_dir))
|
|
except Exception:
|
|
app.logger.exception("Failed to load trace from %s", trace_dir)
|
|
|
|
|
|
@app.route("/trace", methods=["POST"])
|
|
def update_trace():
|
|
data = request.get_json()
|
|
trace_id = data.get("id")
|
|
return_all = data.get("all")
|
|
reset = data.get("reset")
|
|
msg_num = random.randint(1, 10)
|
|
app.logger.info(data)
|
|
log_folder_path = Path(UI_SETTING.trace_folder).absolute()
|
|
if not trace_id:
|
|
return jsonify({"error": "Trace ID is required"}), 400
|
|
trace_id = str(log_folder_path / trace_id)
|
|
|
|
task = _get_or_create_task(trace_id)
|
|
|
|
# Make sure any pending user-interaction requests are visible to the frontend.
|
|
_drain_user_requests_into_messages(task)
|
|
|
|
if task.process is not None and not task.is_alive():
|
|
if not task.messages or task.messages[-1].get("tag") != "END":
|
|
task.messages.append(
|
|
{
|
|
"tag": "END",
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"content": {
|
|
"error_msg": "RD-Agent process has completed.",
|
|
"end_code": task.get_end_code(),
|
|
},
|
|
}
|
|
)
|
|
app.logger.warning(f"Process for {trace_id} has ended.")
|
|
|
|
user_ip = request.remote_addr
|
|
|
|
if reset:
|
|
task.pointers[user_ip] = 0
|
|
|
|
start_pointer = task.pointers[user_ip]
|
|
end_pointer = start_pointer + msg_num
|
|
if end_pointer > len(task.messages) or return_all:
|
|
end_pointer = len(task.messages)
|
|
|
|
returned_msgs = task.messages[start_pointer:end_pointer]
|
|
task.pointers[user_ip] = end_pointer
|
|
if returned_msgs:
|
|
app.logger.info([msg["tag"] for msg in returned_msgs])
|
|
return jsonify(returned_msgs), 200
|
|
|
|
|
|
@app.route("/stdout", methods=["GET"])
|
|
def download_stdout_file():
|
|
trace_id = request.args.get("id", "")
|
|
stdout_path = _resolve_stdout_path(trace_id)
|
|
|
|
if stdout_path is None:
|
|
return jsonify({"error": "Trace ID is required or invalid"}), 400
|
|
if not stdout_path.exists() or not stdout_path.is_file():
|
|
return jsonify({"error": "Stdout file not found"}), 404
|
|
|
|
return send_file(
|
|
stdout_path,
|
|
as_attachment=True,
|
|
download_name=stdout_path.name,
|
|
mimetype="text/plain",
|
|
)
|
|
|
|
|
|
@app.route("/traces", methods=["GET"])
|
|
def list_traces():
|
|
"""Return trace ids that are available for history browsing."""
|
|
|
|
trace_ids = _collect_existing_trace_ids(log_folder_path)
|
|
return jsonify(trace_ids), 200
|
|
|
|
|
|
@app.route("/upload", methods=["POST"])
|
|
def upload_file():
|
|
# 获取请求体中的字段
|
|
global rdagent_processes
|
|
scenario = request.form.get("scenario")
|
|
files = request.files.getlist("files")
|
|
competition = request.form.get("competition")
|
|
loop_n = request.form.get("loops")
|
|
all_duration = request.form.get("all_duration")
|
|
|
|
# scenario = "Data Science Loop"
|
|
if scenario == "Data Science":
|
|
competition = competition[10:] # Eg. MLE-Bench:aerial-cactus-competition
|
|
trace_name = f"{competition}-{randomname.get_name()}"
|
|
else:
|
|
trace_name = randomname.get_name()
|
|
trace_files_path = log_folder_path / "uploads" / scenario / trace_name
|
|
|
|
log_trace_path = (log_folder_path / scenario / trace_name).absolute()
|
|
stdout_path = log_folder_path / scenario / f"{trace_name}.log"
|
|
if not stdout_path.exists():
|
|
stdout_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# save files
|
|
for file in files:
|
|
if file:
|
|
p = (log_folder_path / "uploads" / scenario / trace_name).resolve()
|
|
sanitized_filename = secure_filename(file.filename) # Sanitize filename
|
|
target_path = (p / sanitized_filename).resolve() # Normalize target path
|
|
# Ensure target_path is within the allowed base directory
|
|
if os.path.commonpath([str(target_path), str(p)]) == str(p) and target_path.is_file() == False:
|
|
if not p.exists():
|
|
p.mkdir(parents=True, exist_ok=True)
|
|
file.save(target_path)
|
|
else:
|
|
return jsonify({"error": "Invalid file path"}), 400
|
|
|
|
target_name = None
|
|
kwargs = {}
|
|
loop_n_val = int(loop_n) if loop_n else None
|
|
all_duration_val = f"{all_duration}h" if all_duration else None
|
|
|
|
if scenario == "Finance Data Building":
|
|
target_name = "fin_factor"
|
|
kwargs = {
|
|
"loop_n": loop_n_val,
|
|
"all_duration": all_duration_val,
|
|
"base_features_path": str(trace_files_path),
|
|
}
|
|
if scenario == "Finance Model Implementation":
|
|
target_name = "fin_model"
|
|
kwargs = {
|
|
"loop_n": loop_n_val,
|
|
"all_duration": all_duration_val,
|
|
"base_features_path": str(trace_files_path),
|
|
}
|
|
if scenario == "Finance Whole Pipeline":
|
|
target_name = "fin_quant"
|
|
kwargs = {
|
|
"loop_n": loop_n_val,
|
|
"all_duration": all_duration_val,
|
|
"base_features_path": str(trace_files_path),
|
|
}
|
|
if scenario == "Finance Data Building (Reports)":
|
|
target_name = "fin_factor_report"
|
|
kwargs = {"report_folder": str(trace_files_path), "all_duration": all_duration_val}
|
|
if scenario == "General Model Implementation":
|
|
if len(files) == 0: # files is one link
|
|
rfp = request.form.get("files")[0]
|
|
else: # one file is uploaded
|
|
rfp = str(trace_files_path / files[0].filename)
|
|
target_name = "general_model"
|
|
kwargs = {"report_file_path": rfp}
|
|
if scenario == "Data Science":
|
|
target_name = "data_science"
|
|
kwargs = {"competition": competition, "loop_n": loop_n_val, "timeout": all_duration_val}
|
|
|
|
if target_name is None:
|
|
return jsonify({"error": "Unknown scenario"}), 400
|
|
|
|
app.logger.info(f"Started process for {log_trace_path} with target: {target_name}, kwargs: {kwargs}")
|
|
task = RDAgentTask(
|
|
target_name=target_name,
|
|
kwargs=kwargs,
|
|
stdout_path=str(stdout_path),
|
|
log_trace_path=str(log_trace_path),
|
|
scenario=scenario,
|
|
trace_name=trace_name,
|
|
ui_server_port=app.config["UI_SERVER_PORT"],
|
|
)
|
|
task.start()
|
|
app.logger.warning(f"Task {log_trace_path} started.")
|
|
rdagent_processes[str(log_trace_path)] = task
|
|
return (
|
|
jsonify(
|
|
{
|
|
"id": f"{scenario}/{trace_name}",
|
|
}
|
|
),
|
|
200,
|
|
)
|
|
|
|
|
|
@app.route("/receive", methods=["POST"])
|
|
def receive_msgs():
|
|
try:
|
|
data = request.get_json()
|
|
if not data:
|
|
return jsonify({"error": "No JSON data received"}), 400
|
|
except Exception as e:
|
|
return jsonify({"error": "Internal Server Error"}), 500
|
|
|
|
if isinstance(data, list):
|
|
for d in data:
|
|
task = _get_or_create_task(d["id"])
|
|
task.messages.append(d["msg"])
|
|
else:
|
|
task = _get_or_create_task(data["id"])
|
|
task.messages.append(data["msg"])
|
|
|
|
return jsonify({"status": "success"}), 200
|
|
|
|
|
|
@app.route("/user_interaction/submit", methods=["POST"])
|
|
def submit_user_interaction_response():
|
|
"""Frontend submits a user response; server forwards it to the rdagent subprocess via IPC queue."""
|
|
data = request.get_json(silent=True) or {}
|
|
trace_id = data.get("id")
|
|
payload = data.get("payload")
|
|
|
|
if not trace_id:
|
|
return jsonify({"error": "Trace ID is required"}), 400
|
|
if payload is None:
|
|
return jsonify({"error": "Missing 'payload'"}), 400
|
|
|
|
trace_id = str(log_folder_path / trace_id)
|
|
task = _get_or_create_task(trace_id)
|
|
|
|
try:
|
|
task.user_response_q.put(payload, block=False)
|
|
except Exception as e:
|
|
return jsonify({"error": f"Failed to enqueue user response: {e}"}), 500
|
|
|
|
return jsonify({"status": "success"}), 200
|
|
|
|
|
|
@app.route("/control", methods=["POST"])
|
|
def control_process():
|
|
global rdagent_processes
|
|
data = request.get_json()
|
|
app.logger.info(data)
|
|
if not data or "id" not in data or "action" not in data:
|
|
return jsonify({"error": "Missing 'id' or 'action' in request"}), 400
|
|
|
|
id = str(log_folder_path / data["id"])
|
|
action = data["action"]
|
|
|
|
if action != "stop":
|
|
return jsonify({"error": "Only 'stop' action is supported"}), 400
|
|
|
|
if id not in rdagent_processes or rdagent_processes[id] is None:
|
|
return jsonify({"error": "No running process for given id"}), 400
|
|
|
|
task = rdagent_processes[id]
|
|
|
|
if task.process is None:
|
|
return jsonify({"error": "No running process for given id"}), 400
|
|
|
|
try:
|
|
if task.is_alive():
|
|
task.stop()
|
|
|
|
if not task.messages or task.messages[-1].get("tag") != "END":
|
|
task.messages.append(
|
|
{
|
|
"tag": "END",
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"content": {"error_msg": "RD-Agent process was stopped by user.", "end_code": -1},
|
|
}
|
|
)
|
|
app.logger.warning(f"Process for {id} has been stopped.")
|
|
return jsonify({"status": "stopped"}), 200
|
|
except Exception as e:
|
|
return jsonify({"error": f"Failed to {action} process, {e}"}), 500
|
|
|
|
|
|
@app.route("/test", methods=["GET"])
|
|
def test():
|
|
# return 'Hello, World!'
|
|
msgs = {k: [i["tag"] for i in task.messages] for k, task in rdagent_processes.items()}
|
|
pointers = {k: dict(task.pointers) for k, task in rdagent_processes.items()}
|
|
return jsonify({"msgs": msgs, "pointers": pointers}), 200
|
|
|
|
|
|
@app.route("/", methods=["GET"])
|
|
def index():
|
|
# return 'Hello, World!'
|
|
# return {k: [i["tag"] for i in v] for k, v in msgs_for_frontend.items()}
|
|
return send_from_directory(app.static_folder, "index.html")
|
|
|
|
|
|
@app.route("/<path:fn>", methods=["GET"])
|
|
def server_static_files(fn):
|
|
return send_from_directory(app.static_folder, _normalize_static_request_path(fn))
|
|
|
|
|
|
def main(port: int = 19899):
|
|
app.config["UI_SERVER_PORT"] = port
|
|
_load_existing_traces(log_folder_path)
|
|
app.run(debug=False, host="0.0.0.0", port=port)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(main)
|