Codex
298 lines
11 KiB
Python
298 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import platform
|
|
import queue
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.error import HTTPError
|
|
from urllib.error import URLError
|
|
from urllib.request import Request
|
|
from urllib.request import urlopen
|
|
|
|
|
|
ARTIFACT_NAMES = (
|
|
"lstm_forecaster.json",
|
|
"torch_retrain_guard.json",
|
|
"torch_threshold_calibration.json",
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
repo_root = Path(args.repo_root).resolve()
|
|
runtime_dir = repo_root / "runtime"
|
|
runtime_dir.mkdir(parents=True, exist_ok=True)
|
|
log_path = Path(args.log_file).resolve() if args.log_file else runtime_dir / "windows_training_agent.log"
|
|
|
|
log(log_path, f"TradeBot Windows training agent started for {args.api_base_url}")
|
|
while True:
|
|
try:
|
|
poll_once(args, repo_root, runtime_dir, log_path)
|
|
except Exception as exc: # noqa: BLE001 - agent must keep running.
|
|
log(log_path, f"ERROR: {exc}")
|
|
if args.once:
|
|
break
|
|
time.sleep(max(5, args.poll_seconds))
|
|
|
|
|
|
def poll_once(args: argparse.Namespace, repo_root: Path, runtime_dir: Path, log_path: Path) -> None:
|
|
worker = worker_payload(args, repo_root)
|
|
api_json(args, "/api/training/heartbeat", worker)
|
|
claim = api_json(args, "/api/training/claim", worker)
|
|
if not claim.get("claimed"):
|
|
return
|
|
job = claim.get("job") if isinstance(claim.get("job"), dict) else {}
|
|
job_id = str(job.get("id") or "")
|
|
if not job_id:
|
|
return
|
|
log(log_path, f"Claimed retrain job {job_id}")
|
|
report_progress(args, job_id, "running", "claimed", 2, "Задание получено Windows-agent")
|
|
success = False
|
|
message = ""
|
|
summary: dict[str, Any] = {}
|
|
try:
|
|
run_retrain(args, job_id, job, repo_root, log_path)
|
|
summary = read_json(runtime_dir / "torch_retrain_guard.json")
|
|
report_progress(args, job_id, "running", "uploading", 72, "Обучение завершено, загружаю артефакты")
|
|
for name in ARTIFACT_NAMES:
|
|
path = runtime_dir / name
|
|
if path.is_file():
|
|
upload_artifact(args, job_id, path, log_path)
|
|
success = True
|
|
message = "training completed"
|
|
log(log_path, f"Completed retrain job {job_id}")
|
|
except Exception as exc: # noqa: BLE001 - report failure to the bot.
|
|
message = str(exc)
|
|
log(log_path, f"Job {job_id} failed: {message}")
|
|
finally:
|
|
payload = {"success": success, "message": message, "summary": summary}
|
|
api_json(args, f"/api/training/jobs/{job_id}/complete", payload)
|
|
|
|
|
|
def run_retrain(args: argparse.Namespace, job_id: str, job: dict[str, Any], repo_root: Path, log_path: Path) -> None:
|
|
script = repo_root / "tools" / "run_torch_retrain.ps1"
|
|
if not script.is_file():
|
|
raise RuntimeError(f"retrain script not found: {script}")
|
|
cmd = [
|
|
"powershell.exe",
|
|
"-NoProfile",
|
|
"-ExecutionPolicy",
|
|
"Bypass",
|
|
"-File",
|
|
str(script),
|
|
]
|
|
parameters = job.get("parameters") if isinstance(job.get("parameters"), dict) else {}
|
|
arg_map = {
|
|
"symbols": "-Symbols",
|
|
"limit": "-Limit",
|
|
"lookbacks": "-Lookbacks",
|
|
"architectures": "-Architectures",
|
|
"hidden_sizes": "-HiddenSizes",
|
|
"layers": "-Layers",
|
|
"dropouts": "-Dropouts",
|
|
"epochs": "-Epochs",
|
|
}
|
|
for key, ps_arg in arg_map.items():
|
|
value = parameters.get(key)
|
|
if value not in (None, ""):
|
|
cmd.extend([ps_arg, str(value)])
|
|
log(log_path, "Running retrain: " + " ".join(quote_for_log(part) for part in cmd))
|
|
report_progress(args, job_id, "running", "training", 8, "PyTorch retrain запущен")
|
|
line_count = 0
|
|
output_queue: queue.Queue[str] = queue.Queue()
|
|
|
|
def read_output() -> None:
|
|
assert process.stdout is not None
|
|
for raw_line in process.stdout:
|
|
output_queue.put(raw_line.rstrip())
|
|
|
|
with subprocess.Popen(
|
|
cmd,
|
|
cwd=str(repo_root),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
encoding="utf-8",
|
|
errors="replace",
|
|
) as process:
|
|
reader = threading.Thread(target=read_output, name="training-output-reader", daemon=True)
|
|
reader.start()
|
|
last_report_at = 0.0
|
|
last_message = "PyTorch retrain выполняется"
|
|
while True:
|
|
got_line = False
|
|
try:
|
|
message = output_queue.get(timeout=5)
|
|
got_line = True
|
|
log(log_path, message)
|
|
line_count += 1
|
|
if message:
|
|
last_message = message[-220:]
|
|
except queue.Empty:
|
|
pass
|
|
|
|
progress = min(70, 8 + line_count // 3)
|
|
now = time.monotonic()
|
|
if got_line or now - last_report_at >= 30:
|
|
safe_report_progress(args, job_id, "running", "training", progress, last_message, log_path)
|
|
last_report_at = now
|
|
|
|
if process.poll() is not None and output_queue.empty():
|
|
break
|
|
|
|
reader.join(timeout=2)
|
|
code = process.wait()
|
|
if code != 0:
|
|
raise RuntimeError(f"retrain failed with exit code {code}")
|
|
report_progress(args, job_id, "running", "guard", 70, "Guard завершён, подготавливаю артефакты")
|
|
|
|
|
|
def upload_artifact(args: argparse.Namespace, job_id: str, path: Path, log_path: Path) -> None:
|
|
digest = hashlib.sha256(path.read_bytes()).hexdigest()
|
|
size = path.stat().st_size
|
|
chunk_size = max(64 * 1024, args.chunk_size)
|
|
total = max(1, (size + chunk_size - 1) // chunk_size)
|
|
log(log_path, f"Uploading {path.name}: {size} bytes, {total} chunks")
|
|
with path.open("rb") as source:
|
|
for index in range(total):
|
|
data = source.read(chunk_size)
|
|
payload = {
|
|
"name": path.name,
|
|
"index": index,
|
|
"total": total,
|
|
"sha256": digest,
|
|
"data_base64": base64.b64encode(data).decode("ascii"),
|
|
}
|
|
api_json(args, f"/api/training/jobs/{job_id}/artifacts/chunk", payload, timeout=120)
|
|
if index == 0 or index == total - 1 or index % 10 == 0:
|
|
progress = 72 + int(((index + 1) / total) * 23)
|
|
report_progress(args, job_id, "running", "uploading", progress, f"Загружаю {path.name}: {index + 1}/{total}")
|
|
|
|
|
|
def report_progress(
|
|
args: argparse.Namespace,
|
|
job_id: str,
|
|
status: str,
|
|
phase: str,
|
|
progress_percent: int,
|
|
message: str,
|
|
) -> None:
|
|
api_json(
|
|
args,
|
|
f"/api/training/jobs/{job_id}/progress",
|
|
{
|
|
"status": status,
|
|
"phase": phase,
|
|
"progress_percent": progress_percent,
|
|
"message": message,
|
|
"worker": worker_payload(args, Path(args.repo_root).resolve()),
|
|
},
|
|
)
|
|
|
|
|
|
def safe_report_progress(
|
|
args: argparse.Namespace,
|
|
job_id: str,
|
|
status: str,
|
|
phase: str,
|
|
progress_percent: int,
|
|
message: str,
|
|
log_path: Path,
|
|
) -> None:
|
|
try:
|
|
report_progress(args, job_id, status, phase, progress_percent, message)
|
|
except Exception as exc: # noqa: BLE001 - keep the local training process alive.
|
|
log(log_path, f"Progress report failed: {exc}")
|
|
|
|
|
|
def api_json(args: argparse.Namespace, path: str, payload: dict[str, Any], timeout: int = 30) -> dict[str, Any]:
|
|
url = args.api_base_url.rstrip("/") + path
|
|
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
|
headers = {"Content-Type": "application/json", "Accept": "application/json"}
|
|
token = args.api_auth or os.environ.get("TRADEBOT_API_AUTH", "")
|
|
headers.update(auth_headers(token))
|
|
request = Request(url, data=body, headers=headers, method="POST")
|
|
try:
|
|
with urlopen(request, timeout=timeout) as response:
|
|
text = response.read().decode("utf-8")
|
|
except HTTPError as exc:
|
|
detail = exc.read().decode("utf-8", errors="replace")
|
|
raise RuntimeError(f"HTTP {exc.code} {path}: {detail[:300]}") from exc
|
|
except URLError as exc:
|
|
raise RuntimeError(f"network error {path}: {exc.reason}") from exc
|
|
return json.loads(text) if text.strip() else {}
|
|
|
|
|
|
def auth_headers(token: str) -> dict[str, str]:
|
|
value = token.strip()
|
|
if not value:
|
|
return {}
|
|
headers = {"X-TradeBot-Token": value}
|
|
if value.lower().startswith(("basic ", "bearer ")):
|
|
headers["Authorization"] = value
|
|
elif ":" in value:
|
|
encoded = base64.b64encode(value.encode("utf-8")).decode("ascii")
|
|
headers["Authorization"] = f"Basic {encoded}"
|
|
else:
|
|
headers["Authorization"] = f"Bearer {value}"
|
|
return headers
|
|
|
|
|
|
def worker_payload(args: argparse.Namespace, repo_root: Path) -> dict[str, Any]:
|
|
name = args.worker_name or platform.node() or "Windows training host"
|
|
return {
|
|
"worker_id": args.worker_id or f"{name}:{repo_root}",
|
|
"name": name,
|
|
"path": str(repo_root),
|
|
"version": "1",
|
|
}
|
|
|
|
|
|
def log(path: Path, message: str) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
stamp = datetime.now().astimezone().isoformat(timespec="seconds")
|
|
line = f"[{stamp}] {message}"
|
|
print(line, flush=True)
|
|
with path.open("a", encoding="utf-8") as handle:
|
|
handle.write(line + "\n")
|
|
|
|
|
|
def read_json(path: Path) -> dict[str, Any]:
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
return {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def quote_for_log(value: str) -> str:
|
|
return f'"{value}"' if " " in value else value
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Poll TradeBot for retrain jobs and execute them on Windows.")
|
|
parser.add_argument("--api-base-url", default=os.environ.get("TRADEBOT_API_BASE_URL", "https://tb.kusoft.xyz"))
|
|
parser.add_argument("--api-auth", default=os.environ.get("TRADEBOT_API_AUTH", ""))
|
|
parser.add_argument("--repo-root", default=str(Path(__file__).resolve().parents[1]))
|
|
parser.add_argument("--worker-id", default=os.environ.get("TRADEBOT_TRAINING_WORKER_ID", ""))
|
|
parser.add_argument("--worker-name", default=os.environ.get("TRADEBOT_TRAINING_WORKER_NAME", ""))
|
|
parser.add_argument("--poll-seconds", type=int, default=int(os.environ.get("TRADEBOT_TRAINING_POLL_SECONDS", "60")))
|
|
parser.add_argument("--chunk-size", type=int, default=int(os.environ.get("TRADEBOT_TRAINING_CHUNK_SIZE", str(512 * 1024))))
|
|
parser.add_argument("--log-file", default=os.environ.get("TRADEBOT_TRAINING_LOG", ""))
|
|
parser.add_argument("--once", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|