spiri-sdk/guiTools/spiri_sdk_guitools/sim_drone.py

from loguru import logger
from pathlib import Path
import contextlib
from typing import List
import os
import sh
import subprocess
from nicegui import ui, run, app

import docker
import aiodocker
import asyncio
from spiri_sdk_guitools.video_button import EnableStreamingButton

docker_client = docker.from_env()

from rclpy.executors import ExternalShutdownException
from rclpy.node import Node
import rclpy
import threading


def ros_main() -> None:
    rclpy.init()


app.on_startup(lambda: threading.Thread(target=ros_main).start())


@contextlib.contextmanager
def modified_environ(*remove, **update):
    """
    Temporarily updates the ``os.environ`` dictionary in-place.

    The ``os.environ`` dictionary is updated in-place so that the modification
    is sure to work in all situations.

    :param remove: Environment variables to remove.
    :param update: Dictionary of environment variables and values to add/update.
    """
    env = os.environ
    update = update or {}
    remove = remove or []

    # List of environment variables being updated or removed.
    stomped = (set(update.keys()) | set(remove)) & set(env.keys())
    # Environment variables and values to restore on exit.
    update_after = {k: env[k] for k in stomped}
    # Environment variables and values to remove on exit.
    remove_after = frozenset(k for k in update if k not in env)

    try:
        env.update(update)
        [env.pop(k, None) for k in remove]
        yield
    finally:
        env.update(update_after)
        [env.pop(k) for k in remove_after]


robots = set()


async def container_logs(container, element):
    adocker = aiodocker.Docker()
    with element:
        with ui.scroll_area():
            acontainer = await adocker.containers.get(container.id)
            async for log in acontainer.log(stdout=True, stderr=True, follow=True):
                for line in log.splitlines():
                    ui.label(line)
                # ui.html(conv.convert(bytes(log,'utf-8').decode('utf-8', 'xmlcharrefreplace'), full=False))


class Robot:
    robot_type = "spiri-mu"

    def __init__(self, sysid: int, compose_files: List[Path] | str):
        if sysid > 255 or sysid < 0:
            raise ValueError("sysid must be between 0 and 255")
        self.sysid = int(sysid)
        if isinstance(compose_files, str):
            compose_files = [
                Path(file) for file in compose_files.replace("\n", ",").split(",")
            ]
        self.compose_files = compose_files
        self.processes = []
        self.video_button = None
        robots.add(self)

    async def ui_containers(self, element):
        docker_elements = {}
        container_status = {}
        with element:
            while True:
                # Poll for data that changes
                for container in self.containers():
                    try:
                        health = container.attrs["State"]["Health"]["Status"]
                    except KeyError:
                        health = "Unknown"

                    container_status[container] = (
                        f"{container.name} {container.status} {health}"
                    )
                    if container not in docker_elements:
                        docker_elements[container] = ui.element().classes("w-full")
                        with docker_elements[container]:
                            ui.label().bind_text(container_status, container).classes(
                                "text-lg"
                            )
                            logelement = (
                                ui.expansion("Logs")
                                .style("margin: 10px;")
                                .classes("w-full outline outline-1")
                            )
                            asyncio.create_task(container_logs(container, logelement))
                # Check for containers that have been removed
                removed = set(docker_elements.keys()) - set(self.containers())
                for container in removed:
                    self.robot_ui.remove(docker_elements[container])
                    docker_elements.pop(container)
                await asyncio.sleep(1)

    async def ui_ros(self, element):
        with element:
            node_dummy = Node("_ros2cli_dummy_to_show_topic_list")
            scroll_area = ui.scroll_area()
            with scroll_area:
                while True:
                    scroll_area.clear()
                    for topic in node_dummy.get_topic_names_and_types():
                        ui.label(topic)
                    await asyncio.sleep(10)

    async def ui(self, element):
        adocker = aiodocker.Docker()

        with element:
            self.robot_ui = ui.element().classes("w-full outline p-4")
            with self.robot_ui:
                ui.label(f"{self.robot_type} {self.sysid}").classes("text-2xl")
                ui.label(f"""Sysid: {self.sysid}""")
                ui.button("Start", on_click=self.async_start).classes("m-2")
                ui.button("Stop", on_click=self.async_stop).classes("m-2")
                self.video_button = EnableStreamingButton(sysid=self.sysid).classes(
                    "m-2"
                )

                async def delete_robot():
                    await self.async_stop()
                    robots.remove(self)
                    element.remove(self.robot_ui)

                ui.button("Delete", on_click=delete_robot).classes("m-2")
                with ui.tabs() as tabs:
                    tab_containers = ui.tab("Containers")
                    tab_ros = ui.tab("ROS Topics")
                with ui.tab_panels(tabs, value=tab_containers):
                    tab = ui.tab_panel(tab_containers).classes("w-full")
                    asyncio.create_task(self.ui_containers(tab))
                with ui.tab_panels(tabs, value=tab_ros):
                    tab = ui.tab_panel(tab_ros).classes("w-full")
                    asyncio.create_task(self.ui_ros(tab))

    async def async_stop(self):
        return await run.io_bound(self.stop)

    def stop(self):
        # If video was enabled, and we stop. Button state is wrong.
        # Stop video if it is enabled.
        if isinstance(self.video_button, EnableStreamingButton):
            self.video_button.stop_video()
        # Delete gazebo model
        self.delete_gz_model(sysid=self.sysid)
        # Signal all processes to stop
        for process in self.processes:
            process.terminate()
            process.kill()
        self.processes = []
        for container in self.containers():
            # container.stop()
            container.remove(force=True)

    def containers(self):
        return docker_client.containers.list(
            all=True, filters={"name": f"robot-sim-{self.robot_type}-{self.sysid}"}
        )

    async def async_start(self):
        return await run.io_bound(self.start)

    def start(self):
        """Starts the simulated drone with a given sysid,
        each drone must have it's own unique ID.
        """
        instance = self.sysid - 1
        sysid = self.sysid
        with logger.contextualize(syd_id=sysid):
            env = os.environ
            with modified_environ(
                SERIAL0_PORT=str(int(env["SERIAL0_PORT"]) + 10 * instance),
                MAVROS2_PORT=str(int(env["MAVROS2_PORT"]) + 10 * instance),
                MAVROS1_PORT=str(int(env["MAVROS1_PORT"]) + 10 * instance),
                FDM_PORT_IN=str(int(env["FDM_PORT_IN"]) + 10 * instance),
                GSTREAMER_UDP_PORT=str(int(env["GSTREAMER_UDP_PORT"]) + 10 * instance),
                SITL_PORT=str(int(env["SITL_PORT"]) + 10 * instance),
                INSTANCE=str(instance),
                DRONE_SYS_ID=str(self.sysid),
            ):
                self.spawn_gz_model(self.sysid)
                logger.info("Starting drone stack, this may take some time")
                for compose_file in self.compose_files:
                    if not isinstance(compose_file, Path):
                        compose_file = Path(compose_file)
                    if not compose_file.exists():
                        raise FileNotFoundError(f"File {compose_file} does not exist")
                    args = [
                        "docker-compose",
                        "--profile",
                        "uav-sim",
                        "-p",
                        f"robot-sim-{self.robot_type}-{sysid}",
                        "-f",
                        compose_file.as_posix(),
                        "up",
                    ]
                    command = " ".join(args)

                    logger.info(f"Starting drone stack with command: {command}")
                    docker_stack = subprocess.Popen(
                        args,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                    )

    @staticmethod
    def spawn_gz_model(sysid):
        logger.info("")
        env = os.environ
        GSTREAMER_UDP_PORT = env["GSTREAMER_UDP_PORT"]
        FDM_PORT_IN = env["FDM_PORT_IN"]
        WORLD_NAME = env["WORLD_NAME"]
        DRONE_MODEL = env["DRONE_MODEL"]
        XACRO_CMD = [
            "xacro",
            f"gstreamer_udp_port:={GSTREAMER_UDP_PORT}",
            f"fdm_port_in:={FDM_PORT_IN}",
            "model.xacro.sdf",
            "-o",
            "model.sdf",
        ]
        # This path is breaking if this drone_model folder doesnt exist!
        # TODO: fix this model path for minimal code maintenance
        ROS2_CMD = f"ros2 run ros_gz_sim create -world {WORLD_NAME} -file /ardupilot_gazebo/models/{DRONE_MODEL}/model.sdf -name spiri-{sysid} -x {sysid - 1}  -y 0 -z 0.195"

        xacro_proc = subprocess.Popen(
            XACRO_CMD,
            cwd=f"/ardupilot_gazebo/models/{DRONE_MODEL}",
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        ros2_gz_create_proc = subprocess.Popen(
            ROS2_CMD.split(),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        out, err = ros2_gz_create_proc.communicate(timeout=15)
        ros2_gz_create_proc.kill()
        return

    @staticmethod
    def delete_gz_model(sysid):
        env = os.environ
        WORLD_NAME = env["WORLD_NAME"]
        # http://osrf-distributions.s3.amazonaws.com/gazebo/api/7.1.0/classgazebo_1_1physics_1_1Entity.html
        ENTITY_TYPE_MODEL = 0x00000002
        REQUEST_ARG = f"name: 'spiri-{sysid}' type: {ENTITY_TYPE_MODEL}"
        GZ_SERVICE_CMD = [
            "gz",
            "service",
            "-s",
            f"/world/{WORLD_NAME}/remove",
            "--reqtype",
            "gz.msgs.Entity",
            "--reptype",
            "gz.msgs.Boolean",
            "--timeout",
            "5000",
            "--req",
            REQUEST_ARG,
        ]
        remove_entity_proc = subprocess.Popen(
            GZ_SERVICE_CMD,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        out, err = remove_entity_proc.communicate(timeout=15)
        remove_entity_proc.kill()