dora-ulite6/dora_ulite6/main.py

"""The dora event loop: accept requests, publish what the arm threads produced.

Nothing here touches the arm. The loop reads the `request_id` dora puts in a
message's metadata, hands the work to the control thread, and publishes the
answer under the same id.

Two services do not take the ordinary road. `get_state` runs here, because it
never blocks and is wanted exactly when the control thread is busy. `stop` is
handed to the stopper thread, because it must not queue behind the motion it
aborts and must not be run here either — the SDK spends up to three seconds on
it, and a router cannot go away that long.
"""

import argparse
import os

import pyarrow as pa
from dora import Node

from dora_ulite6 import payload, services
from dora_ulite6.config import Config
from dora_ulite6.control import Control, Message, ServiceCall, Telemetry
from dora_ulite6.log import log

# The metadata key dora's service pattern is built on.
REQUEST_ID = "request_id"


def main() -> None:
    """Connect to the dataflow, start the arm threads, and run until dora stops us."""
    args = _parse_args()
    config = Config.from_env()

    # The node comes first, and the arm afterwards. Starting the arm threads
    # first would let INIT_MODE drive a real robot before this process had
    # joined the dataflow at all — and if joining then failed, it would have
    # moved with nobody watching. Connecting first also lets the arm be told
    # whether this is a restart.
    #
    # Passing the name lets the node also run standalone as a dynamic node; when
    # dora spawns it, the id from the dataflow wins.
    node = Node(args.name)

    control = Control(config, restarted=node.is_restart())
    control.start()

    try:
        for event in node:
            if event["type"] == "ERROR":
                raise RuntimeError(event["error"])
            if event["type"] != "INPUT":
                continue
            _handle(node, control, event)
    finally:
        if not control.shutdown():
            # The arm has been told to stop, but the control thread is still
            # inside a blocking SDK call and may never come out. Leave rather
            # than hang until dora loses patience and kills us.
            os._exit(0)


def _handle(node: Node, control: Control, event: dict) -> None:
    """Route one input event."""
    event_id = event["id"]
    if event_id == "tick":
        _publish(node, control)
    elif event_id == "request":
        _accept_request(control, event)
    elif event_id == "jog":
        _accept_jog(control, event)
    else:
        log(f"ignoring unexpected input `{event_id}`")


def _publish(node: Node, control: Control) -> None:
    """Send everything the arm threads produced since the last tick."""
    for state in control.take_state_changes():
        node.send_output("robot_state", pa.array([state]))
    for message in control.take_messages():
        node.send_output(message.output, pa.array([message.fields]), metadata=message.metadata)
    telemetry = control.take_telemetry()
    if telemetry is not None:
        _send_telemetry(node, telemetry)


def _send_telemetry(node: Node, telemetry: Telemetry) -> None:
    node.send_output(
        "joint_pose",
        pa.array(telemetry.joints, type=pa.float32()),
        metadata={"units": "deg", "timestamp_ns": telemetry.timestamp_ns},
    )
    node.send_output(
        "tcp_pose",
        pa.array(telemetry.tip.as_list(), type=pa.float32()),
        metadata={
            "units_position": "mm",
            "units_rotation": "deg",
            "frame": "tool_tip",
            "timestamp_ns": telemetry.timestamp_ns,
        },
    )


def _accept_request(control: Control, event: dict) -> None:
    """Take one service request, or refuse it under its own request id.

    A request the loop cannot parse is answered here rather than queued: a
    caller waiting on a request_id must get exactly one response, never silence.
    A request it *can* parse is queued even before the arm has been reached,
    because `recover` and `get_state` are precisely what a caller sends then.

    `wait` decides only *when* the one response is sent — at once, or when the
    arm stops. It is read here rather than in the service handler because it is
    about the protocol, not about what the arm does.
    """
    request_id = event["metadata"].get(REQUEST_ID, "")
    try:
        row = payload.read_row(event["value"])
        name = payload.name_of(row, services.SERVICE, services.NAMES)
        wait = payload.flag(row, "wait", True)
    except payload.PayloadError as error:
        _refuse(control, request_id, str(error))
        return
    if name == services.STOP:
        control.request_stop(request_id)
        return
    if name in services.IMMEDIATE:
        control.call_now(request_id, row)
        return
    call = ServiceCall(request_id=request_id, row=row, wait=wait)
    if call.moves_the_arm:
        # Motions go through their own door: one preempts whatever is already
        # moving, rather than queueing behind it.
        control.start_motion(call)
        return
    control.submit(call)


def _accept_jog(control: Control, event: dict) -> None:
    """Take one velocity command.

    A topic rather than a service: a jog is a stream that expires on the
    controller's own TTL, so a sender keeps sending and nobody waits for an
    answer. It still drives the arm, so it preempts a motion and reports
    `moving` the same way one does — an arm that is travelling must never look
    idle.
    """
    try:
        row = payload.read_row(event["value"])
        velocity = payload.numbers(row, "velocity", count=6)
    except payload.PayloadError as error:
        log(f"ignoring a malformed jog: {error}")
        return
    control.start_jog(velocity)


def _refuse(control: Control, request_id: str, message: str) -> None:
    """Answer a request the loop rejected, without troubling the arm threads."""
    control.post(
        Message(
            output="response",
            fields={"ok": False, "service": "", "message": message},
            metadata={REQUEST_ID: request_id},
        )
    )


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Dora node for a UFactory Lite6 arm.")
    parser.add_argument(
        "--name",
        default="ulite6",
        help="The name of the node in the dataflow, used when running as a dynamic node.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    main()