463 lines
16 KiB
Python
463 lines
16 KiB
Python
|
"""RT-OSS Node Agent — connects to RIA Hub and dispatches work to local hardware.
|
|||
|
|
|
|||
|
|
The agent runs on any machine with an SDR attached and connects **outbound** to
|
|||
|
|
RIA Hub. No inbound ports need to be opened on the user's machine, and the
|
|||
|
|
connection works identically through NAT, corporate firewalls, or a Pi on a
|
|||
|
|
cellular link.
|
|||
|
|
|
|||
|
|
Usage::
|
|||
|
|
|
|||
|
|
ria-agent \\
|
|||
|
|
--hub https://riahub.company.com \\
|
|||
|
|
--key <api-key> \\
|
|||
|
|
--name lab-bench-1 \\
|
|||
|
|
[--device plutosdr] \\
|
|||
|
|
[--insecure]
|
|||
|
|
|
|||
|
|
The agent:
|
|||
|
|
1. Registers with RIA Hub and receives a ``node_id``.
|
|||
|
|
2. Sends a heartbeat every 30 s so the hub knows it is online.
|
|||
|
|
3. Long-polls ``GET /orchestrator/nodes/{id}/commands`` (30 s timeout).
|
|||
|
|
4. Executes received campaigns via :class:`ria_toolkit_oss.orchestration.executor.CampaignExecutor`.
|
|||
|
|
5. Uploads recordings to the hub via chunked POST, keeping each request
|
|||
|
|
under 50 MB so it passes through Cloudflare without needing the bypass
|
|||
|
|
subdomain.
|
|||
|
|
6. Deregisters cleanly on SIGINT / SIGTERM.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import math
|
|||
|
|
import os
|
|||
|
|
import signal
|
|||
|
|
import sys
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
import uuid
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("ria_agent")
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# Tuneable constants
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
_HEARTBEAT_INTERVAL = 30 # seconds between heartbeats
|
|||
|
|
_POLL_TIMEOUT = 30 # server-side long-poll duration
|
|||
|
|
_POLL_CLIENT_TIMEOUT = 40 # client read timeout — slightly longer than server
|
|||
|
|
_RECONNECT_PAUSE = 5 # seconds to wait after a poll error before retrying
|
|||
|
|
_CHUNK_SIZE = 50 * 1024 * 1024 # 50 MB — well below Cloudflare's 100 MB limit
|
|||
|
|
_DIRECT_THRESHOLD = 90 * 1024 * 1024 # files above this use chunked upload
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# Agent
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
class NodeAgent:
|
|||
|
|
"""Outbound-connecting agent that bridges RIA Hub to local SDR hardware.
|
|||
|
|
|
|||
|
|
All network I/O is initiated by the agent (outbound). RIA Hub never opens
|
|||
|
|
a connection back to the agent's machine.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
hub_url: str,
|
|||
|
|
api_key: str,
|
|||
|
|
name: str,
|
|||
|
|
sdr_device: str = "unknown",
|
|||
|
|
insecure: bool = False,
|
|||
|
|
) -> None:
|
|||
|
|
self.hub_url = hub_url.rstrip("/")
|
|||
|
|
self.api_key = api_key
|
|||
|
|
self.name = name
|
|||
|
|
self.sdr_device = sdr_device
|
|||
|
|
self.insecure = insecure
|
|||
|
|
|
|||
|
|
self.node_id: str | None = None
|
|||
|
|
self._stop = threading.Event()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
import ria_toolkit_oss
|
|||
|
|
|
|||
|
|
self._ria_version: str = getattr(ria_toolkit_oss, "__version__", "unknown")
|
|||
|
|
except Exception:
|
|||
|
|
self._ria_version = "unknown"
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Public entry point
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def run(self) -> None:
|
|||
|
|
"""Register, start the heartbeat thread, and enter the command loop.
|
|||
|
|
|
|||
|
|
Blocks until SIGINT or SIGTERM is received.
|
|||
|
|
"""
|
|||
|
|
self._register()
|
|||
|
|
|
|||
|
|
def _shutdown(sig: int, _frame: Any) -> None:
|
|||
|
|
logger.info("Shutdown signal received — stopping agent")
|
|||
|
|
self._stop.set()
|
|||
|
|
|
|||
|
|
signal.signal(signal.SIGINT, _shutdown)
|
|||
|
|
signal.signal(signal.SIGTERM, _shutdown)
|
|||
|
|
|
|||
|
|
hb = threading.Thread(target=self._heartbeat_loop, daemon=True, name="ria-agent-heartbeat")
|
|||
|
|
hb.start()
|
|||
|
|
|
|||
|
|
logger.info("Agent %r online (node_id=%s, hub=%s)", self.name, self.node_id, self.hub_url)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
self._command_loop()
|
|||
|
|
finally:
|
|||
|
|
self._stop.set()
|
|||
|
|
self._deregister()
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Registration
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _register(self) -> None:
|
|||
|
|
resp = self._post(
|
|||
|
|
"/orchestrator/nodes/register",
|
|||
|
|
json={
|
|||
|
|
"name": self.name,
|
|||
|
|
"sdr_device": self.sdr_device,
|
|||
|
|
"ria_toolkit_version": self._ria_version,
|
|||
|
|
"capabilities": ["inference", "campaign"],
|
|||
|
|
},
|
|||
|
|
timeout=15,
|
|||
|
|
)
|
|||
|
|
resp.raise_for_status()
|
|||
|
|
self.node_id = resp.json()["node_id"]
|
|||
|
|
logger.info("Registered as %r (node_id=%s)", self.name, self.node_id)
|
|||
|
|
|
|||
|
|
def _deregister(self) -> None:
|
|||
|
|
if not self.node_id:
|
|||
|
|
return
|
|||
|
|
try:
|
|||
|
|
self._delete(f"/orchestrator/nodes/{self.node_id}", timeout=10)
|
|||
|
|
logger.info("Deregistered %s", self.node_id)
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.debug("Deregister failed (ignored on shutdown): %s", exc)
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Heartbeat thread
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _heartbeat_loop(self) -> None:
|
|||
|
|
while not self._stop.wait(_HEARTBEAT_INTERVAL):
|
|||
|
|
try:
|
|||
|
|
resp = self._post(f"/orchestrator/nodes/{self.node_id}/heartbeat", timeout=10)
|
|||
|
|
if resp.status_code == 404:
|
|||
|
|
logger.warning("Heartbeat got 404 — hub lost registration, re-registering")
|
|||
|
|
self._register()
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.warning("Heartbeat failed: %s", exc)
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Command poll loop
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _command_loop(self) -> None:
|
|||
|
|
while not self._stop.is_set():
|
|||
|
|
try:
|
|||
|
|
resp = self._get(
|
|||
|
|
f"/orchestrator/nodes/{self.node_id}/commands",
|
|||
|
|
timeout=_POLL_CLIENT_TIMEOUT,
|
|||
|
|
)
|
|||
|
|
if resp.status_code == 204:
|
|||
|
|
# No command within the timeout window — loop immediately.
|
|||
|
|
continue
|
|||
|
|
if resp.status_code == 404:
|
|||
|
|
logger.warning("Command poll got 404 — re-registering")
|
|||
|
|
self._register()
|
|||
|
|
continue
|
|||
|
|
resp.raise_for_status()
|
|||
|
|
cmd = resp.json()
|
|||
|
|
logger.info("Received command: %s", cmd.get("command"))
|
|||
|
|
self._dispatch(cmd)
|
|||
|
|
except Exception as exc:
|
|||
|
|
if not self._stop.is_set():
|
|||
|
|
logger.warning("Command poll error: %s — retrying in %ds", exc, _RECONNECT_PAUSE)
|
|||
|
|
time.sleep(_RECONNECT_PAUSE)
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Command dispatch
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _dispatch(self, cmd: dict) -> None:
|
|||
|
|
command = cmd.get("command")
|
|||
|
|
if command == "run_campaign":
|
|||
|
|
campaign_id: str = cmd.get("campaign_id") or str(uuid.uuid4())
|
|||
|
|
config_dict: dict = cmd.get("payload") or {}
|
|||
|
|
threading.Thread(
|
|||
|
|
target=self._run_campaign,
|
|||
|
|
args=(campaign_id, config_dict),
|
|||
|
|
daemon=True,
|
|||
|
|
name=f"campaign-{campaign_id[:8]}",
|
|||
|
|
).start()
|
|||
|
|
else:
|
|||
|
|
logger.warning("Unknown command %r — ignored", command)
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Campaign execution
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _run_campaign(self, campaign_id: str, config_dict: dict) -> None:
|
|||
|
|
try:
|
|||
|
|
from ria_toolkit_oss.orchestration.campaign import CampaignConfig
|
|||
|
|
from ria_toolkit_oss.orchestration.executor import CampaignExecutor
|
|||
|
|
except ImportError as exc:
|
|||
|
|
logger.error(
|
|||
|
|
"Campaign %s cannot start — ria_toolkit_oss not fully installed: %s",
|
|||
|
|
campaign_id[:8],
|
|||
|
|
exc,
|
|||
|
|
)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
logger.info("Campaign %s starting", campaign_id[:8])
|
|||
|
|
try:
|
|||
|
|
config = CampaignConfig.from_dict(config_dict)
|
|||
|
|
executor = CampaignExecutor(config)
|
|||
|
|
result = executor.run()
|
|||
|
|
logger.info("Campaign %s completed — uploading recordings", campaign_id[:8])
|
|||
|
|
self._upload_recordings(campaign_id, config, result)
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.error("Campaign %s failed: %s", campaign_id[:8], exc)
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Recording upload (chunked for large files)
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _upload_recordings(self, campaign_id: str, config: Any, result: Any) -> None:
|
|||
|
|
output_repo: str | None = getattr(getattr(config, "output", None), "repo", None)
|
|||
|
|
if not output_repo or "/" not in output_repo:
|
|||
|
|
logger.warning("Campaign %s: no output.repo — skipping upload", campaign_id[:8])
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
repo_owner, repo_name = output_repo.split("/", 1)
|
|||
|
|
base_url = f"{self.hub_url}/datasets/upload"
|
|||
|
|
steps = getattr(result, "steps", None) or []
|
|||
|
|
|
|||
|
|
for step in steps:
|
|||
|
|
output_path: str | None = getattr(step, "output_path", None)
|
|||
|
|
if not output_path:
|
|||
|
|
continue
|
|||
|
|
device_id: str = getattr(step, "transmitter_id", "") or ""
|
|||
|
|
for fpath in _sigmf_files(output_path):
|
|||
|
|
filename = os.path.basename(fpath)
|
|||
|
|
metadata = {
|
|||
|
|
"filename": filename,
|
|||
|
|
"repo_owner": repo_owner,
|
|||
|
|
"repo_name": repo_name,
|
|||
|
|
"device_id": device_id,
|
|||
|
|
"campaign_id": campaign_id,
|
|||
|
|
}
|
|||
|
|
try:
|
|||
|
|
resp_data = self._upload_file(base_url, fpath, metadata)
|
|||
|
|
logger.info(
|
|||
|
|
"Campaign %s: uploaded %s (oid=%s)",
|
|||
|
|
campaign_id[:8],
|
|||
|
|
filename,
|
|||
|
|
resp_data.get("oid", "?"),
|
|||
|
|
)
|
|||
|
|
except Exception as exc:
|
|||
|
|
logger.warning("Campaign %s: upload of %s failed: %s", campaign_id[:8], filename, exc)
|
|||
|
|
|
|||
|
|
def _upload_file(self, base_url: str, file_path: str, metadata: dict) -> dict:
|
|||
|
|
"""Upload *file_path*, choosing chunked or direct path based on file size."""
|
|||
|
|
import requests as _requests
|
|||
|
|
|
|||
|
|
size = os.path.getsize(file_path)
|
|||
|
|
filename = os.path.basename(file_path)
|
|||
|
|
headers = {"X-API-Key": self.api_key}
|
|||
|
|
verify = not self.insecure
|
|||
|
|
|
|||
|
|
# Small files: single POST (unchanged endpoint, no assembly needed server-side).
|
|||
|
|
if size <= _DIRECT_THRESHOLD:
|
|||
|
|
with open(file_path, "rb") as fh:
|
|||
|
|
resp = _requests.post(
|
|||
|
|
base_url,
|
|||
|
|
headers=headers,
|
|||
|
|
files={"file": (filename, fh)},
|
|||
|
|
data=metadata,
|
|||
|
|
timeout=300,
|
|||
|
|
verify=verify,
|
|||
|
|
)
|
|||
|
|
resp.raise_for_status()
|
|||
|
|
return resp.json()
|
|||
|
|
|
|||
|
|
# Large files: chunked upload — each request is ≤ 50 MB.
|
|||
|
|
total_chunks = math.ceil(size / _CHUNK_SIZE)
|
|||
|
|
upload_id = str(uuid.uuid4())
|
|||
|
|
chunk_url = base_url + "/chunk"
|
|||
|
|
|
|||
|
|
logger.info(
|
|||
|
|
"Chunked upload: %s (%d bytes, %d × %d MB chunks)",
|
|||
|
|
filename,
|
|||
|
|
size,
|
|||
|
|
total_chunks,
|
|||
|
|
_CHUNK_SIZE // (1024 * 1024),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
resp_data: dict = {}
|
|||
|
|
with open(file_path, "rb") as fh:
|
|||
|
|
for i in range(total_chunks):
|
|||
|
|
chunk = fh.read(_CHUNK_SIZE)
|
|||
|
|
resp = _requests.post(
|
|||
|
|
chunk_url,
|
|||
|
|
headers=headers,
|
|||
|
|
files={"file": (filename, chunk, "application/octet-stream")},
|
|||
|
|
data={
|
|||
|
|
**metadata,
|
|||
|
|
"upload_id": upload_id,
|
|||
|
|
"chunk_index": i,
|
|||
|
|
"total_chunks": total_chunks,
|
|||
|
|
},
|
|||
|
|
timeout=120,
|
|||
|
|
verify=verify,
|
|||
|
|
)
|
|||
|
|
if not resp.ok:
|
|||
|
|
raise RuntimeError(
|
|||
|
|
f"Chunk {i + 1}/{total_chunks} failed: " f"HTTP {resp.status_code}: {resp.text[:300]}"
|
|||
|
|
)
|
|||
|
|
resp_data = resp.json()
|
|||
|
|
logger.debug("Chunk %d/%d uploaded", i + 1, total_chunks)
|
|||
|
|
|
|||
|
|
return resp_data
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# HTTP helpers
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _get(self, path: str, **kwargs: Any):
|
|||
|
|
import requests as _requests
|
|||
|
|
|
|||
|
|
return _requests.get(
|
|||
|
|
f"{self.hub_url}{path}",
|
|||
|
|
headers={"X-API-Key": self.api_key},
|
|||
|
|
verify=not self.insecure,
|
|||
|
|
**kwargs,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _post(self, path: str, **kwargs: Any):
|
|||
|
|
import requests as _requests
|
|||
|
|
|
|||
|
|
return _requests.post(
|
|||
|
|
f"{self.hub_url}{path}",
|
|||
|
|
headers={"X-API-Key": self.api_key},
|
|||
|
|
verify=not self.insecure,
|
|||
|
|
**kwargs,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _delete(self, path: str, **kwargs: Any):
|
|||
|
|
import requests as _requests
|
|||
|
|
|
|||
|
|
return _requests.delete(
|
|||
|
|
f"{self.hub_url}{path}",
|
|||
|
|
headers={"X-API-Key": self.api_key},
|
|||
|
|
verify=not self.insecure,
|
|||
|
|
**kwargs,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# Helpers
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _sigmf_files(data_path: str) -> list[str]:
|
|||
|
|
"""Return paths to both SigMF files (.sigmf-data and .sigmf-meta) for a recording."""
|
|||
|
|
candidates = [data_path]
|
|||
|
|
if data_path.endswith(".sigmf-data"):
|
|||
|
|
candidates.append(data_path[: -len(".sigmf-data")] + ".sigmf-meta")
|
|||
|
|
return [p for p in candidates if os.path.exists(p)]
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# CLI entry point
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main() -> None:
|
|||
|
|
import argparse
|
|||
|
|
|
|||
|
|
parser = argparse.ArgumentParser(
|
|||
|
|
prog="ria-agent",
|
|||
|
|
description=(
|
|||
|
|
"RT-OSS Node Agent — connects outbound to RIA Hub and executes "
|
|||
|
|
"campaigns / inference on local SDR hardware."
|
|||
|
|
),
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--hub",
|
|||
|
|
required=True,
|
|||
|
|
metavar="URL",
|
|||
|
|
help="RIA Hub base URL, e.g. https://riahub.company.com",
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--key",
|
|||
|
|
required=True,
|
|||
|
|
metavar="API_KEY",
|
|||
|
|
help="Shared API key (must match [wac] API_KEY in the hub's app.ini)",
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--name",
|
|||
|
|
required=True,
|
|||
|
|
metavar="NAME",
|
|||
|
|
help='Human-readable name shown in the Target Node dropdown, e.g. "lab-bench-1"',
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--device",
|
|||
|
|
default="unknown",
|
|||
|
|
metavar="SDR",
|
|||
|
|
help=(
|
|||
|
|
"SDR device type reported to the hub (informational only). "
|
|||
|
|
"Examples: plutosdr, usrp_b210, rtlsdr, mock. Default: unknown"
|
|||
|
|
),
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--insecure",
|
|||
|
|
action="store_true",
|
|||
|
|
help="Disable TLS certificate verification (dev/self-signed certs only)",
|
|||
|
|
)
|
|||
|
|
parser.add_argument(
|
|||
|
|
"--log-level",
|
|||
|
|
default="INFO",
|
|||
|
|
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
|
|||
|
|
help="Logging verbosity (default: INFO)",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
logging.basicConfig(
|
|||
|
|
level=getattr(logging, args.log_level),
|
|||
|
|
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
|||
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|||
|
|
stream=sys.stderr,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Warn loudly if --insecure is used outside of development.
|
|||
|
|
if args.insecure:
|
|||
|
|
logger.warning(
|
|||
|
|
"--insecure disables TLS certificate verification. "
|
|||
|
|
"Only use this for local development with self-signed certs."
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
agent = NodeAgent(
|
|||
|
|
hub_url=args.hub,
|
|||
|
|
api_key=args.key,
|
|||
|
|
name=args.name,
|
|||
|
|
sdr_device=args.device,
|
|||
|
|
insecure=args.insecure,
|
|||
|
|
)
|
|||
|
|
agent.run()
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|