114 lines
3.9 KiB
Python
Executable File
114 lines
3.9 KiB
Python
Executable File
#!/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()
|