Add script to test actor/context/job/flow create and flow dag
Signed-off-by: Lee Smet <lee.smet@hotmail.com>
This commit is contained in:
Lee Smet
361
scripts/jsonrpc_demo.py
Normal file
361
scripts/jsonrpc_demo.py
Normal file
@@ -0,0 +1,361 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Demo script for HeroCoordinator JSON-RPC API.
|
||||
- Creates an actor
|
||||
- Verifies by loading the actor
|
||||
- Creates a context with the actor as admin/reader/executor
|
||||
- Creates three jobs with dependencies
|
||||
- Creates a flow referencing those jobs
|
||||
- Fetches and prints the flow DAG
|
||||
|
||||
Usage:
|
||||
COORDINATOR_URL=http://127.0.0.1:9652 python3 scripts/jsonrpc_demo.py
|
||||
Defaults to http://127.0.0.1:9652 if COORDINATOR_URL is not set.
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import sys
|
||||
from urllib import request, error
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
JSONRPC_VERSION = "2.0"
|
||||
|
||||
class JsonRpcClient:
|
||||
def __init__(self, url: str):
|
||||
self.url = url.rstrip("/")
|
||||
self._id = 0
|
||||
|
||||
def call(self, method: str, params: Dict[str, Any]) -> Any:
|
||||
self._id += 1
|
||||
payload = {
|
||||
"jsonrpc": JSONRPC_VERSION,
|
||||
"id": self._id,
|
||||
"method": method,
|
||||
"params": params,
|
||||
}
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = request.Request(self.url, data=data, headers={"Content-Type": "application/json"})
|
||||
try:
|
||||
with request.urlopen(req) as resp:
|
||||
body = resp.read()
|
||||
except error.HTTPError as e:
|
||||
try:
|
||||
details = e.read().decode("utf-8", "ignore")
|
||||
except Exception:
|
||||
details = ""
|
||||
raise RuntimeError(f"HTTP error {e.code}: {details}") from e
|
||||
except error.URLError as e:
|
||||
raise RuntimeError(f"URL error: {e.reason}") from e
|
||||
|
||||
try:
|
||||
obj = json.loads(body.decode("utf-8"))
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Invalid JSON response: {body!r}") from e
|
||||
|
||||
# JSON-RPC single response expected
|
||||
if isinstance(obj, list):
|
||||
raise RuntimeError("Batch responses are not supported in this demo")
|
||||
|
||||
if obj.get("error"):
|
||||
raise RuntimeError(f"RPC error: {json.dumps(obj['error'])}")
|
||||
|
||||
return obj.get("result")
|
||||
|
||||
def print_header(title: str):
|
||||
print("\n" + "=" * 80)
|
||||
print(title)
|
||||
print("=" * 80)
|
||||
|
||||
def pretty_print(obj: Any):
|
||||
print(json.dumps(obj, indent=2, sort_keys=True))
|
||||
|
||||
def summarize_dag(dag: Dict[str, Any]):
|
||||
print_header("Flow DAG Summary")
|
||||
flow_id = dag.get("flow_id")
|
||||
caller_id = dag.get("caller_id")
|
||||
context_id = dag.get("context_id")
|
||||
print(f"flow_id={flow_id} caller_id={caller_id} context_id={context_id}")
|
||||
edges: List[Tuple[int, int]] = dag.get("edges", [])
|
||||
roots: List[int] = dag.get("roots", [])
|
||||
leaves: List[int] = dag.get("leaves", [])
|
||||
levels: List[List[int]] = dag.get("levels", [])
|
||||
nodes: Dict[str, Any] = dag.get("nodes", {})
|
||||
|
||||
print("Edges:")
|
||||
for a, b in edges:
|
||||
print(f" {a} -> {b}")
|
||||
|
||||
print(f"Roots: {roots}")
|
||||
print(f"Leaves: {leaves}")
|
||||
print("Levels:")
|
||||
for i, lvl in enumerate(levels):
|
||||
print(f" L{i}: {lvl}")
|
||||
|
||||
# Show nodes and their dependencies (from JobSummary)
|
||||
print("Nodes:")
|
||||
for k, v in nodes.items():
|
||||
depends = v.get("depends", [])
|
||||
prerequisites = v.get("prerequisites", [])
|
||||
stype = v.get("script_type")
|
||||
print(f" Job {k}: depends={depends} prerequisites={prerequisites} script_type={stype}")
|
||||
|
||||
def assert_edges(edges: List[Tuple[int, int]], required: List[Tuple[int, int]]):
|
||||
edge_set = {(int(a), int(b)) for a, b in edges}
|
||||
missing = [e for e in required if e not in edge_set]
|
||||
if missing:
|
||||
raise AssertionError(f"Missing expected edges in DAG: {missing}; got={sorted(edge_set)}")
|
||||
|
||||
def main():
|
||||
url = os.getenv("COORDINATOR_URL", "http://127.0.0.1:9652")
|
||||
client = JsonRpcClient(url)
|
||||
|
||||
# Deterministic demo IDs; change if collisions happen
|
||||
actor_id = 1001
|
||||
context_id = 1 # Redis DB indices are 0-15; keep <= 15
|
||||
job_a = 3001
|
||||
job_b = 3002
|
||||
job_c = 3003
|
||||
job_d = 3004
|
||||
job_e = 3005
|
||||
job_f = 3006
|
||||
job_g = 3007
|
||||
job_h = 3008
|
||||
job_i = 3009
|
||||
flow_id = 4001
|
||||
|
||||
print_header("actor.create")
|
||||
actor = client.call("actor.create", {
|
||||
"actor": {
|
||||
"id": actor_id,
|
||||
"pubkey": "demo-pubkey",
|
||||
"address": ["127.0.0.1"]
|
||||
}
|
||||
})
|
||||
pretty_print(actor)
|
||||
|
||||
print_header("actor.load")
|
||||
actor_loaded = client.call("actor.load", {"id": actor_id})
|
||||
pretty_print(actor_loaded)
|
||||
|
||||
print_header("context.create")
|
||||
context = client.call("context.create", {
|
||||
"context": {
|
||||
"id": context_id,
|
||||
"admins": [actor_id],
|
||||
"readers": [actor_id],
|
||||
"executors": [actor_id]
|
||||
}
|
||||
})
|
||||
pretty_print(context)
|
||||
|
||||
print_header("job.create - A (root)")
|
||||
jobA = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_a,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('A')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": []
|
||||
}
|
||||
})
|
||||
pretty_print(jobA)
|
||||
|
||||
print_header("job.create - B (root)")
|
||||
jobB = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_b,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('B')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": []
|
||||
}
|
||||
})
|
||||
pretty_print(jobB)
|
||||
|
||||
print_header("job.create - C (depends on A and B)")
|
||||
jobC = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_c,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('C')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": [job_a, job_b]
|
||||
}
|
||||
})
|
||||
pretty_print(jobC)
|
||||
|
||||
print_header("job.create - D (depends on A)")
|
||||
jobD = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_d,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('D')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": [job_a]
|
||||
}
|
||||
})
|
||||
pretty_print(jobD)
|
||||
|
||||
print_header("job.create - E (depends on B)")
|
||||
jobE = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_e,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('E')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": [job_b]
|
||||
}
|
||||
})
|
||||
pretty_print(jobE)
|
||||
|
||||
print_header("job.create - F (depends on C and D)")
|
||||
jobF = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_f,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('F')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": [job_c, job_d]
|
||||
}
|
||||
})
|
||||
pretty_print(jobF)
|
||||
|
||||
print_header("job.create - G (depends on C and E)")
|
||||
jobG = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_g,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('G')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": [job_c, job_e]
|
||||
}
|
||||
})
|
||||
pretty_print(jobG)
|
||||
|
||||
print_header("job.create - H (leaf; depends on F and G)")
|
||||
jobH = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_h,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('H')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": [job_f, job_g]
|
||||
}
|
||||
})
|
||||
pretty_print(jobH)
|
||||
|
||||
print_header("job.create - I (leaf; depends on F and G)")
|
||||
jobI = client.call("job.create", {
|
||||
"context_id": context_id,
|
||||
"job": {
|
||||
"id": job_i,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"script": "print('I')",
|
||||
"script_type": "Python",
|
||||
"timeout": 30,
|
||||
"retries": 0,
|
||||
"env_vars": {},
|
||||
"prerequisites": [],
|
||||
"depends": [job_f, job_g]
|
||||
}
|
||||
})
|
||||
pretty_print(jobI)
|
||||
|
||||
print_header("flow.create")
|
||||
flow = client.call("flow.create", {
|
||||
"context_id": context_id,
|
||||
"flow": {
|
||||
"id": flow_id,
|
||||
"caller_id": actor_id,
|
||||
"context_id": context_id,
|
||||
"jobs": [job_a, job_b, job_c, job_d, job_e, job_f, job_g, job_h, job_i],
|
||||
"env_vars": {}
|
||||
}
|
||||
})
|
||||
pretty_print(flow)
|
||||
|
||||
print_header("flow.dag")
|
||||
dag = client.call("flow.dag", {"context_id": context_id, "id": flow_id})
|
||||
summarize_dag(dag)
|
||||
|
||||
# Validate roots and leaves
|
||||
got_roots = list(map(int, dag.get("roots", [])))
|
||||
if got_roots != sorted([job_a, job_b]):
|
||||
print("WARNING: Unexpected roots:", got_roots, file=sys.stderr)
|
||||
|
||||
got_leaves = {int(x) for x in dag.get("leaves", [])}
|
||||
expected_leaves = {job_h, job_i}
|
||||
if got_leaves != expected_leaves:
|
||||
print("WARNING: Unexpected leaves:", got_leaves, "expected:", expected_leaves, file=sys.stderr)
|
||||
|
||||
# Check edges reflect the expanded DAG
|
||||
expected_edges = [
|
||||
(job_a, job_c), (job_b, job_c),
|
||||
(job_a, job_d), (job_b, job_e),
|
||||
(job_c, job_f), (job_d, job_f),
|
||||
(job_c, job_g), (job_e, job_g),
|
||||
(job_f, job_h), (job_g, job_h),
|
||||
(job_f, job_i), (job_g, job_i),
|
||||
]
|
||||
try:
|
||||
assert_edges(dag.get("edges", []), expected_edges)
|
||||
print("DAG edges contain expected dependencies:", expected_edges)
|
||||
except AssertionError as e:
|
||||
print("WARNING:", e, file=sys.stderr)
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except Exception as e:
|
||||
print_header("Error")
|
||||
print(str(e))
|
||||
sys.exit(1)
|
||||
Reference in New Issue
Block a user