first

domain-api/deploy/multi-region/simulate_cluster_node.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+API_ROOT = Path(__file__).resolve().parents[2]
+if str(API_ROOT) not in sys.path:
+    sys.path.insert(0, str(API_ROOT))
+
+from app.core.db import get_db
+from app.services.cluster_runtime_service import ensure_runtime_schema, register_node_heartbeat
+
+
+def _metadata_from_args(args: argparse.Namespace) -> dict:
+    metadata = {
+        "service": "cluster-simulator",
+        "simulated": True,
+        "phase": args.phase or "",
+        "detail": args.detail or "",
+    }
+    if args.job_id:
+        metadata["job_id"] = args.job_id
+    if args.job_code:
+        metadata["job_code"] = args.job_code
+    if args.cycle_token:
+        metadata["cycle_token"] = args.cycle_token
+    if args.active_threads is not None:
+        metadata["active_threads"] = args.active_threads
+    if args.available_proxy_count is not None:
+        metadata["available_proxy_count"] = args.available_proxy_count
+    if args.metadata_json:
+        try:
+            custom = json.loads(args.metadata_json)
+            if isinstance(custom, dict):
+                metadata.update(custom)
+        except Exception as exc:
+            raise SystemExit(f"invalid --metadata-json: {exc}") from exc
+    return metadata
+
+
+def _delete_node(node_code: str) -> None:
+    with get_db() as conn:
+        with conn.cursor() as cur:
+            cur.execute("DELETE FROM detect_worker_nodes WHERE node_code = %s", (str(node_code),))
+        conn.commit()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Simulate a cluster node heartbeat for multi-region rehearsal.")
+    parser.add_argument("--node-code", required=True)
+    parser.add_argument("--region", required=True)
+    parser.add_argument("--role", required=True, choices=("control", "worker"))
+    parser.add_argument("--status", default="online")
+    parser.add_argument("--current-load", type=int, default=0)
+    parser.add_argument("--phase", default="")
+    parser.add_argument("--detail", default="")
+    parser.add_argument("--job-id", type=int, default=0)
+    parser.add_argument("--job-code", default="")
+    parser.add_argument("--cycle-token", default="")
+    parser.add_argument("--active-threads", type=int, default=None)
+    parser.add_argument("--available-proxy-count", type=int, default=None)
+    parser.add_argument("--metadata-json", default="")
+    parser.add_argument("--interval", type=int, default=20)
+    parser.add_argument("--iterations", type=int, default=0, help="0 means run forever")
+    parser.add_argument("--clear", action="store_true", help="delete the node record and exit")
+    args = parser.parse_args()
+
+    ensure_runtime_schema()
+    if args.clear:
+        _delete_node(args.node_code)
+        print(f"cleared simulated node: {args.node_code}")
+        return
+
+    metadata = _metadata_from_args(args)
+    interval = max(5, int(args.interval or 20))
+    iterations = max(0, int(args.iterations or 0))
+    current = 0
+    while True:
+        current += 1
+        register_node_heartbeat(
+            node_code=args.node_code,
+            region=args.region,
+            role=args.role,
+            status=args.status,
+            current_load=max(0, int(args.current_load or 0)),
+            metadata=metadata,
+        )
+        print(
+            json.dumps(
+                {
+                    "node_code": args.node_code,
+                    "region": args.region,
+                    "role": args.role,
+                    "status": args.status,
+                    "current_load": max(0, int(args.current_load or 0)),
+                    "iteration": current,
+                    "metadata": metadata,
+                },
+                ensure_ascii=False,
+            ),
+            flush=True,
+        )
+        if iterations and current >= iterations:
+            break
+        time.sleep(interval)
+
+
+if __name__ == "__main__":
+    main()