from __future__ import annotations import argparse import base64 import hashlib import json import os import platform import subprocess import sys 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}") success = False message = "" summary: dict[str, Any] = {} try: run_retrain(job, repo_root, log_path) summary = read_json(runtime_dir / "torch_retrain_guard.json") 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(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)) with subprocess.Popen( cmd, cwd=str(repo_root), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", errors="replace", ) as process: assert process.stdout is not None for line in process.stdout: log(log_path, line.rstrip()) code = process.wait() if code != 0: raise RuntimeError(f"retrain failed with exit code {code}") 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) 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()