Codex
332 lines
12 KiB
Python
332 lines
12 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
|
||
started_at = time.monotonic()
|
||
last_output_at = started_at
|
||
last_message = "PyTorch retrain выполняется"
|
||
while True:
|
||
got_line = False
|
||
try:
|
||
message = output_queue.get(timeout=5)
|
||
got_line = True
|
||
last_output_at = time.monotonic()
|
||
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:
|
||
report_message = last_message
|
||
if not got_line:
|
||
report_message = training_heartbeat_message(now, started_at, last_output_at, last_message)
|
||
safe_report_progress(args, job_id, "running", "training", progress, report_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 training_heartbeat_message(now: float, started_at: float, last_output_at: float, last_message: str) -> str:
|
||
elapsed = format_duration(now - started_at)
|
||
idle_seconds = max(0.0, now - last_output_at)
|
||
if idle_seconds >= 45:
|
||
return (
|
||
f"PyTorch обучает модель: процесс активен {elapsed}; "
|
||
f"последний лог {format_duration(idle_seconds)} назад: {last_message[:140]}"
|
||
)
|
||
return last_message or f"PyTorch обучает модель: процесс активен {elapsed}"
|
||
|
||
|
||
def format_duration(seconds: float) -> str:
|
||
total_seconds = max(0, int(seconds))
|
||
minutes, seconds_part = divmod(total_seconds, 60)
|
||
hours, minutes_part = divmod(minutes, 60)
|
||
if hours:
|
||
return f"{hours}ч {minutes_part}м"
|
||
if minutes:
|
||
return f"{minutes}м {seconds_part}с"
|
||
return f"{seconds_part}с"
|
||
|
||
|
||
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:
|
||
last_error: Exception | None = None
|
||
for attempt in range(1, 4):
|
||
try:
|
||
report_progress(args, job_id, status, phase, progress_percent, message)
|
||
return
|
||
except Exception as exc: # noqa: BLE001 - keep the local training process alive.
|
||
last_error = exc
|
||
if attempt < 3:
|
||
time.sleep(attempt * 2)
|
||
log(log_path, f"Temporary progress upload error; training continues: {last_error}")
|
||
|
||
|
||
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()
|