439 lines
16 KiB
Python
439 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Supervisor flow demo for HeroCoordinator.
|
|
|
|
This script:
|
|
- Optionally pre-registers a Python runner on the target Supervisor over Mycelium using an admin secret (--admin-secret). If the flag is not set, this step is skipped.
|
|
- Creates an actor
|
|
- Creates a context granting the actor admin/reader/executor privileges
|
|
- Registers a Runner in the context targeting a Supervisor reachable via Mycelium (by public key or IP)
|
|
- Creates simple Python jobs (text jobs) with a small dependency chain
|
|
- Creates a flow referencing those jobs
|
|
- Starts the flow and polls until it finishes (or errors)
|
|
|
|
Transport: JSON-RPC over HTTP to the Coordinator (default COORDINATOR_URL=http://127.0.0.1:9652).
|
|
|
|
Example usage:
|
|
COORDINATOR_URL=http://127.0.0.1:9652 python3 scripts/supervisor_flow_demo.py --dst-ip 2001:db8::1 [--secret your-secret]
|
|
COORDINATOR_URL=http://127.0.0.1:9652 python3 scripts/supervisor_flow_demo.py --dst-pk bb39b4a3a4efd70f3e05e37887677e02efbda14681d0acd3882bc0f754792c32 [--secret your-secret]
|
|
|
|
Notes:
|
|
- Exactly one of --dst-ip or --dst-pk must be provided.
|
|
- Runner.topic defaults to "supervisor.rpc" (see main.rs).
|
|
- The router auto-discovers contexts and will deliver job.run messages to the supervisor.
|
|
- Mycelium URL is read from MYCELIUM_URL (default http://127.0.0.1:8990).
|
|
- supervisor.register_runner uses static name="python" and queue="python".
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import base64
|
|
import os
|
|
import sys
|
|
import time
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
from urllib import request, error
|
|
|
|
JSONRPC_VERSION = "2.0"
|
|
|
|
|
|
def env_url() -> str:
|
|
return os.getenv("COORDINATOR_URL", "http://127.0.0.1:9652").rstrip("/")
|
|
|
|
def env_mycelium_url() -> str:
|
|
return os.getenv("MYCELIUM_URL", "http://127.0.0.1:8990").rstrip("/")
|
|
|
|
|
|
class JsonRpcClient:
|
|
def __init__(self, url: str):
|
|
self.url = url
|
|
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
|
|
|
|
if isinstance(obj, list):
|
|
raise RuntimeError("Batch responses are not supported")
|
|
|
|
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(obj: Any):
|
|
print(json.dumps(obj, indent=2, sort_keys=True))
|
|
|
|
def mycelium_register_runner(
|
|
myc: "JsonRpcClient",
|
|
dst_pk: Optional[str],
|
|
dst_ip: Optional[str],
|
|
topic: str,
|
|
admin_secret: str,
|
|
name: str = "python",
|
|
queue: str = "python",
|
|
timeout: int = 15,
|
|
) -> Any:
|
|
"""
|
|
Send supervisor.register_runner over Mycelium using pushMessage and wait for the reply.
|
|
- myc: JsonRpcClient for the Mycelium API (MYCELIUM_URL)
|
|
- dst_pk/dst_ip: destination on the overlay; one of them must be provided
|
|
- topic: message topic (defaults to supervisor.rpc from args)
|
|
- admin_secret: supervisor admin secret to authorize the registration
|
|
- name/queue: static identifiers for the python runner on the supervisor
|
|
- timeout: seconds to wait for a reply
|
|
Returns the JSON-RPC 'result' from the supervisor or raises on error/timeout.
|
|
"""
|
|
envelope = {
|
|
"jsonrpc": JSONRPC_VERSION,
|
|
"id": 1,
|
|
"method": "register_runner",
|
|
"params": [{"secret": admin_secret, "name": name, "queue": queue}],
|
|
}
|
|
payload_b64 = base64.b64encode(json.dumps(envelope).encode("utf-8")).decode("ascii")
|
|
topic_b64 = base64.b64encode(topic.encode("utf-8")).decode("ascii")
|
|
|
|
if dst_pk:
|
|
dst = {"pk": dst_pk}
|
|
elif dst_ip:
|
|
dst = {"ip": dst_ip}
|
|
else:
|
|
raise RuntimeError("Either dst_pk or dst_ip must be provided for Mycelium destination")
|
|
|
|
params = {
|
|
"message": {"dst": dst, "topic": topic_b64, "payload": payload_b64},
|
|
}
|
|
resp = myc.call("pushMessage", params)
|
|
time.sleep(15)
|
|
|
|
# Expect an InboundMessage with a payload if a reply was received
|
|
# if isinstance(resp, dict) and "payload" in resp:
|
|
# try:
|
|
# reply = json.loads(base64.b64decode(resp["payload"]).decode("utf-8"))
|
|
# except Exception as e:
|
|
# raise RuntimeError(f"Invalid supervisor reply payload: {e}")
|
|
# if isinstance(reply, dict) and reply.get("error"):
|
|
# raise RuntimeError(f"Supervisor register_runner error: {json.dumps(reply['error'])}")
|
|
# return reply.get("result")
|
|
#
|
|
# raise RuntimeError("No reply received from supervisor for register_runner (timeout)")
|
|
|
|
|
|
def try_create_or_load(client: JsonRpcClient, create_method: str, create_params: Dict[str, Any],
|
|
load_method: str, load_params: Dict[str, Any]) -> Any:
|
|
"""Attempt a create; if it fails due to existence, try load."""
|
|
try:
|
|
return client.call(create_method, create_params)
|
|
except RuntimeError as e:
|
|
msg = str(e)
|
|
# Server maps AlreadyExists to StorageError, we don't have a structured error code here.
|
|
if "Already exists" in msg or "Storage Error" in msg or "Invalid params" in msg:
|
|
# Fall back to load
|
|
return client.call(load_method, load_params)
|
|
raise
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(description="Create actor/context/runner/jobs/flow; start and wait until completion.")
|
|
group = p.add_mutually_exclusive_group(required=True)
|
|
group.add_argument("--dst-ip", help="Supervisor Mycelium IP address (IPv4 or IPv6)")
|
|
group.add_argument("--dst-pk", help="Supervisor public key (64-hex)")
|
|
|
|
p.add_argument("--context-id", type=int, default=2, help="Context id (Redis DB index; 0-15). Default: 2")
|
|
p.add_argument("--actor-id", type=int, default=11001, help="Actor id. Default: 11001")
|
|
p.add_argument("--runner-id", type=int, default=12001, help="Runner id. Default: 12001")
|
|
p.add_argument("--flow-id", type=int, default=13001, help="Flow id. Default: 13001")
|
|
p.add_argument("--base-job-id", type=int, default=20000, help="Base job id for first job; subsequent jobs increment. Default: 20000")
|
|
p.add_argument("--jobs", type=int, default=3, help="Number of jobs to create (>=1). Forms a simple chain. Default: 3")
|
|
p.add_argument("--timeout-secs", type=int, default=60, help="Per-job timeout seconds. Default: 60")
|
|
p.add_argument("--retries", type=int, default=0, help="Per-job retries (0-255). Default: 0")
|
|
p.add_argument(
|
|
"--script-type",
|
|
choices=["Python", "V", "Osis", "Sal"],
|
|
default="Python",
|
|
help="ScriptType for jobs/runner. Default: Python"
|
|
)
|
|
p.add_argument("--topic", default="supervisor.rpc", help="Supervisor topic. Default: supervisor.rpc")
|
|
p.add_argument("--secret", help="Optional supervisor secret used for authenticated supervisor calls")
|
|
p.add_argument("--admin-secret", help="Supervisor admin secret to pre-register a Python runner over Mycelium. If omitted, pre-registration is skipped.")
|
|
p.add_argument("--poll-interval", type=float, default=2.0, help="Flow poll interval seconds. Default: 2.0")
|
|
p.add_argument("--poll-timeout", type=int, default=600, help="Max seconds to wait for flow completion. Default: 600")
|
|
return p.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
if args.jobs < 1:
|
|
print("ERROR: --jobs must be >= 1", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
url = env_url()
|
|
client = JsonRpcClient(url)
|
|
|
|
mycelium_url = env_mycelium_url()
|
|
mycelium_client = JsonRpcClient(mycelium_url) if getattr(args, "admin_secret", None) else None
|
|
|
|
actor_id = int(args.actor_id)
|
|
context_id = int(args.context_id)
|
|
runner_id = int(args.runner_id)
|
|
flow_id = int(args.flow_id)
|
|
base_job_id = int(args.base_job_id)
|
|
script_type = args.script_type
|
|
timeout = int(args.timeout_secs)
|
|
retries = int(args.retries)
|
|
topic = args.topic
|
|
|
|
# 1) Actor
|
|
print_header("actor.create (or load)")
|
|
actor = try_create_or_load(
|
|
client,
|
|
"actor.create",
|
|
{
|
|
"actor": {
|
|
"id": actor_id,
|
|
"pubkey": "demo-pubkey",
|
|
"address": ["127.0.0.1"],
|
|
}
|
|
},
|
|
"actor.load",
|
|
{"id": actor_id},
|
|
)
|
|
pretty(actor)
|
|
|
|
# 2) Context
|
|
print_header("context.create (or load)")
|
|
context = try_create_or_load(
|
|
client,
|
|
"context.create",
|
|
{
|
|
"context": {
|
|
"id": context_id,
|
|
"admins": [actor_id],
|
|
"readers": [actor_id],
|
|
"executors": [actor_id],
|
|
}
|
|
},
|
|
"context.load",
|
|
{"id": context_id},
|
|
)
|
|
pretty(context)
|
|
|
|
# 3) Runner in this context
|
|
# Router picks pubkey if non-empty, else IP address.
|
|
# However, RunnerCreate requires both fields; we fill both and control routing via pubkey empty/non-empty.
|
|
runner_pubkey = args.dst_pk if args.dst_pk else ""
|
|
runner_address = args.dst_ip if args.dst_ip else "127.0.0.1"
|
|
|
|
# Optional: pre-register a Python runner on the Supervisor over Mycelium using an admin secret
|
|
if getattr(args, "admin_secret", None):
|
|
print_header("supervisor.register_runner (pre-register via Mycelium)")
|
|
try:
|
|
mycelium_result = mycelium_register_runner(
|
|
mycelium_client,
|
|
args.dst_pk if args.dst_pk else None,
|
|
args.dst_ip if args.dst_ip else None,
|
|
topic,
|
|
args.admin_secret,
|
|
name="Python",
|
|
queue="Python",
|
|
timeout=15,
|
|
)
|
|
print("Supervisor register_runner ->", mycelium_result)
|
|
except Exception as e:
|
|
print(f"ERROR: Supervisor pre-registration failed: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
print_header("runner.create (or load)")
|
|
# runner.load requires both context_id and id
|
|
try:
|
|
runner_payload = {
|
|
"id": runner_id,
|
|
"pubkey": runner_pubkey,
|
|
"address": runner_address,
|
|
"topic": topic,
|
|
"script_type": script_type,
|
|
"local": False,
|
|
}
|
|
# Optional supervisor secret used by router for authenticated supervisor calls
|
|
if getattr(args, "secret", None):
|
|
runner_payload["secret"] = args.secret
|
|
runner = client.call("runner.create", {
|
|
"context_id": context_id,
|
|
"runner": runner_payload
|
|
})
|
|
except RuntimeError as e:
|
|
msg = str(e)
|
|
if "Already exists" in msg or "Storage Error" in msg or "Invalid params" in msg:
|
|
runner = client.call("runner.load", {"context_id": context_id, "id": runner_id})
|
|
else:
|
|
raise
|
|
pretty(runner)
|
|
|
|
# 4) Jobs
|
|
# Build a simple chain: J0 (root), J1 depends on J0, J2 depends on J1, ... up to N-1
|
|
job_ids: List[int] = []
|
|
for i in range(args.jobs):
|
|
jid = base_job_id + i
|
|
depends = [] if i == 0 else [base_job_id + (i - 1)]
|
|
job_payload = {
|
|
"id": jid,
|
|
"caller_id": actor_id,
|
|
"context_id": context_id,
|
|
"script": f"print('Job {i} running')",
|
|
"script_type": script_type,
|
|
"timeout": timeout,
|
|
"retries": retries,
|
|
"env_vars": {},
|
|
"prerequisites": [],
|
|
"depends": depends,
|
|
}
|
|
print_header(f"job.create - {jid} {'(root)' if not depends else f'(depends on {depends})'}")
|
|
try:
|
|
job = client.call("job.create", {
|
|
"context_id": context_id,
|
|
"job": job_payload
|
|
})
|
|
except RuntimeError as e:
|
|
msg = str(e)
|
|
if "Already exists" in msg or "Storage Error" in msg or "Invalid params" in msg:
|
|
job = client.call("job.load", {
|
|
"context_id": context_id,
|
|
"caller_id": actor_id,
|
|
"id": jid
|
|
})
|
|
else:
|
|
raise
|
|
pretty(job)
|
|
job_ids.append(jid)
|
|
|
|
# 5) Flow
|
|
print_header("flow.create (or load)")
|
|
try:
|
|
flow = client.call("flow.create", {
|
|
"context_id": context_id,
|
|
"flow": {
|
|
"id": flow_id,
|
|
"caller_id": actor_id,
|
|
"context_id": context_id,
|
|
"jobs": job_ids,
|
|
"env_vars": {}
|
|
}
|
|
})
|
|
except RuntimeError as e:
|
|
msg = str(e)
|
|
if "Already exists" in msg or "Storage Error" in msg or "Invalid params" in msg:
|
|
flow = client.call("flow.load", {"context_id": context_id, "id": flow_id})
|
|
else:
|
|
raise
|
|
pretty(flow)
|
|
|
|
# Optional: show DAG
|
|
try:
|
|
print_header("flow.dag")
|
|
dag = client.call("flow.dag", {"context_id": context_id, "id": flow_id})
|
|
pretty(dag)
|
|
except Exception as e:
|
|
print(f"WARN: flow.dag failed: {e}", file=sys.stderr)
|
|
|
|
# 6) Start flow (idempotent; returns bool whether scheduler started)
|
|
print_header("flow.start")
|
|
started = client.call("flow.start", {"context_id": context_id, "id": flow_id})
|
|
print(f"flow.start -> {started}")
|
|
|
|
# 7) Poll until Finished or Error (or timeout)
|
|
print_header("Polling flow.load until completion")
|
|
t0 = time.time()
|
|
status = None
|
|
last_status_print = 0.0
|
|
poll_count = 0
|
|
while True:
|
|
poll_count += 1
|
|
flow = client.call("flow.load", {"context_id": context_id, "id": flow_id})
|
|
status = flow.get("status")
|
|
now = time.time()
|
|
if now - last_status_print >= max(1.0, float(args.poll_interval)):
|
|
print(f"[{int(now - t0)}s] flow.status = {status}")
|
|
last_status_print = now
|
|
|
|
# Every 5th poll, print the current flow DAG
|
|
if (poll_count % 5) == 0:
|
|
try:
|
|
print_header("flow.dag (periodic)")
|
|
dag = client.call("flow.dag", {"context_id": context_id, "id": flow_id})
|
|
pretty(dag)
|
|
except Exception as e:
|
|
print(f"WARN: periodic flow.dag failed: {e}", file=sys.stderr)
|
|
|
|
if status in ("Finished", "Error"):
|
|
break
|
|
if (now - t0) > args.poll_timeout:
|
|
print(f"ERROR: Flow did not complete within {args.poll_timeout}s (status={status})", file=sys.stderr)
|
|
break
|
|
|
|
time.sleep(float(args.poll_interval))
|
|
|
|
# 8) Final summary: job statuses
|
|
print_header("Final job statuses")
|
|
for jid in job_ids:
|
|
try:
|
|
j = client.call("job.load", {
|
|
"context_id": context_id,
|
|
"caller_id": actor_id,
|
|
"id": jid
|
|
})
|
|
print(f"Job {jid}: status={j.get('status')} result={j.get('result')}")
|
|
except Exception as e:
|
|
print(f"Job {jid}: load failed: {e}", file=sys.stderr)
|
|
|
|
# Exit code
|
|
if status == "Finished":
|
|
print_header("Result")
|
|
print("Flow finished successfully.")
|
|
sys.exit(0)
|
|
else:
|
|
print_header("Result")
|
|
print(f"Flow ended with status={status}")
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
print("\nInterrupted.")
|
|
sys.exit(130)
|
|
except Exception as e:
|
|
print_header("Error")
|
|
print(str(e))
|
|
sys.exit(1)
|