|
| 1 | +#!/usr/bin/env python3 |
| 2 | +import argparse |
| 3 | +import json |
| 4 | +import os |
| 5 | +import sys |
| 6 | +import time |
| 7 | +from pathlib import Path |
| 8 | + |
| 9 | + |
| 10 | +def parse_args() -> argparse.Namespace: |
| 11 | + parser = argparse.ArgumentParser(description="Run one DSR agent worker process") |
| 12 | + parser.add_argument("--agent-name", required=True) |
| 13 | + parser.add_argument("--agent-id", required=True, type=int) |
| 14 | + parser.add_argument("--domain-id", required=True, type=int) |
| 15 | + parser.add_argument("--same-host", required=True, choices=("true", "false")) |
| 16 | + parser.add_argument("--graph-file", default="") |
| 17 | + parser.add_argument("--artifacts-dir", required=True) |
| 18 | + parser.add_argument("--local-attr", required=True) |
| 19 | + parser.add_argument("--local-value", required=True) |
| 20 | + parser.add_argument("--remote-attr", required=True) |
| 21 | + parser.add_argument("--remote-value", required=True) |
| 22 | + parser.add_argument("--startup-delay", default=0.0, type=float) |
| 23 | + parser.add_argument("--sync-timeout", default=30.0, type=float) |
| 24 | + parser.add_argument("--hold-seconds", default=0.0, type=float) |
| 25 | + return parser.parse_args() |
| 26 | + |
| 27 | + |
| 28 | +def wait_for(predicate, timeout_s: float, interval_s: float = 0.1, error: str = "timeout"): |
| 29 | + deadline = time.monotonic() + timeout_s |
| 30 | + while time.monotonic() < deadline: |
| 31 | + value = predicate() |
| 32 | + if value: |
| 33 | + return value |
| 34 | + time.sleep(interval_s) |
| 35 | + raise TimeoutError(error) |
| 36 | + |
| 37 | + |
| 38 | +def read_root_attr(graph, attr_name: str): |
| 39 | + root = graph.get_node("root") |
| 40 | + if root is None: |
| 41 | + return None |
| 42 | + if attr_name not in root.attrs: |
| 43 | + return None |
| 44 | + return root.attrs[attr_name].value |
| 45 | + |
| 46 | + |
| 47 | +def main() -> int: |
| 48 | + args = parse_args() |
| 49 | + artifacts_dir = Path(args.artifacts_dir) |
| 50 | + artifacts_dir.mkdir(parents=True, exist_ok=True) |
| 51 | + result_path = artifacts_dir / f"{args.agent_name}.json" |
| 52 | + |
| 53 | + build_python_wrapper = Path(__file__).resolve().parents[2] / "build" / "python-wrapper" |
| 54 | + sys.path.insert(0, str(build_python_wrapper)) |
| 55 | + |
| 56 | + import pydsr |
| 57 | + |
| 58 | + time.sleep(args.startup_delay) |
| 59 | + |
| 60 | + graph = pydsr.DSRGraph( |
| 61 | + 0, |
| 62 | + args.agent_name, |
| 63 | + args.agent_id, |
| 64 | + args.graph_file, |
| 65 | + args.same_host == "true", |
| 66 | + args.domain_id, |
| 67 | + ) |
| 68 | + |
| 69 | + result = { |
| 70 | + "agent_name": args.agent_name, |
| 71 | + "agent_id": args.agent_id, |
| 72 | + "domain_id": args.domain_id, |
| 73 | + "same_host": args.same_host == "true", |
| 74 | + "graph_file_loaded": bool(args.graph_file), |
| 75 | + } |
| 76 | + |
| 77 | + try: |
| 78 | + initial_nodes = wait_for( |
| 79 | + lambda: len(graph.get_nodes()) if graph.get_node("root") is not None else 0, |
| 80 | + timeout_s=args.sync_timeout, |
| 81 | + error="graph root never became available", |
| 82 | + ) |
| 83 | + result["initial_node_count"] = initial_nodes |
| 84 | + |
| 85 | + root = wait_for( |
| 86 | + lambda: graph.get_node("root"), |
| 87 | + timeout_s=args.sync_timeout, |
| 88 | + error="root node not available", |
| 89 | + ) |
| 90 | + root.attrs[args.local_attr] = pydsr.Attribute(args.local_value) |
| 91 | + update_ok = graph.update_node(root) |
| 92 | + if not update_ok: |
| 93 | + raise RuntimeError(f"failed to update root with {args.local_attr}") |
| 94 | + |
| 95 | + observed_remote = wait_for( |
| 96 | + lambda: read_root_attr(graph, args.remote_attr), |
| 97 | + timeout_s=args.sync_timeout, |
| 98 | + error=f"remote attribute {args.remote_attr} not observed", |
| 99 | + ) |
| 100 | + if observed_remote != args.remote_value: |
| 101 | + raise RuntimeError( |
| 102 | + f"unexpected value for {args.remote_attr}: {observed_remote!r} != {args.remote_value!r}" |
| 103 | + ) |
| 104 | + |
| 105 | + final_root = graph.get_node("root") |
| 106 | + result["final_node_count"] = len(graph.get_nodes()) |
| 107 | + result["local_attr_value"] = final_root.attrs[args.local_attr].value |
| 108 | + result["remote_attr_value"] = final_root.attrs[args.remote_attr].value |
| 109 | + if args.hold_seconds > 0: |
| 110 | + time.sleep(args.hold_seconds) |
| 111 | + result["status"] = "ok" |
| 112 | + except Exception as exc: |
| 113 | + result["status"] = "error" |
| 114 | + result["error"] = str(exc) |
| 115 | + finally: |
| 116 | + result_path.write_text(json.dumps(result, indent=2), encoding="utf-8") |
| 117 | + |
| 118 | + return 0 if result["status"] == "ok" else 1 |
| 119 | + |
| 120 | + |
| 121 | +if __name__ == "__main__": |
| 122 | + raise SystemExit(main()) |
0 commit comments