From e18ff535325af374c9d03e8e4d4ad1caf3052d0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=93=D1=80=D0=B0?= =?UTF-8?q?=D0=B1=D0=B0=D1=80=D1=8C?= Date: Thu, 21 May 2026 11:41:54 +1000 Subject: [PATCH] Add run command for launching robot controller and Webots simulator --- cobot/cli.py | 5 + cobot/commands/run.py | 402 ++++++++++++++++++++++++++++++++++++++++++ cobot/tui.py | 92 +++++++++- 3 files changed, 498 insertions(+), 1 deletion(-) create mode 100644 cobot/commands/run.py diff --git a/cobot/cli.py b/cobot/cli.py index 2e3d383..f7ff4fb 100644 --- a/cobot/cli.py +++ b/cobot/cli.py @@ -6,6 +6,7 @@ from cobot.commands import docker_setup as cmd_docker_setup from cobot.commands import doc_setup as cmd_doc_setup from cobot.commands import local_setup as cmd_local_setup from cobot.commands import robot_setup as cmd_robot_setup +from cobot.commands import run as cmd_run from cobot.commands import setup as cmd_setup from cobot.commands import update as cmd_update @@ -19,6 +20,9 @@ _GROUPS = [ ("doc-setup", "deploy or stop the MkDocs documentation server"), ("robot-setup", "configure cobot-setting.yaml interactively"), ]), + ("Run", [ + ("run", "launch the robot controller or Webots simulator (local or Docker)"), + ]), ("Management", [ ("update", "pull latest changes from the remote git branch and reinstall cobot"), ("delete", "remove the project, Docker images, containers, and optionally ROS2"), @@ -80,5 +84,6 @@ def _register_commands(subparsers): cmd_docker_setup.register(subparsers) cmd_doc_setup.register(subparsers) cmd_robot_setup.register(subparsers) + cmd_run.register(subparsers) cmd_update.register(subparsers) cmd_delete.register(subparsers) diff --git a/cobot/commands/run.py b/cobot/commands/run.py new file mode 100644 index 0000000..9cd002f --- /dev/null +++ b/cobot/commands/run.py @@ -0,0 +1,402 @@ +from __future__ import annotations + +import argparse +import os +import signal +import shutil +import subprocess +from pathlib import Path +from typing import Callable, List, Optional + +from textual.app import App + +from cobot.tui import SCREEN_CSS, LogScreen, PickScreen, RunScreen + +_PROJECT_DIR = Path(__file__).parent.parent.parent +_CONFIG_PATH = _PROJECT_DIR / "cobot-setting.yaml" +_INSTALL_DIR = _PROJECT_DIR / "install" +_JAZZY_DIR = Path("/opt/ros/jazzy") +_CONFIG_IN_CONTAINER = "/ros2_ws/cobot-setting.yaml" + +_CONTAINER_CONTROLLER = "lwc-controller" +_CONTAINER_WEBOTS = "lwc-webots" + +# Candidates checked in order; for controller the webots image is a valid fallback +_CONTROLLER_IMAGES = [ + "lwc-local:ros-iiwa7-jazzy", + "evilfisru/lwc:iiwa-jazzy", + "evilfisru/lwc:iiwa-jazzy-dev", + "lwc-local:ros-iiwa7-webots-jazzy", + "evilfisru/lwc:webots-jazzy", + "evilfisru/lwc:webots-jazzy-dev", +] +_WEBOTS_IMAGES = [ + "lwc-local:ros-iiwa7-webots-jazzy", + "evilfisru/lwc:webots-jazzy", + "evilfisru/lwc:webots-jazzy-dev", +] + + +# --------------------------------------------------------------------------- +# Small utilities +# --------------------------------------------------------------------------- + +class _Ask(App[Optional[str]]): + CSS = SCREEN_CSS + + def __init__(self, step: str, question: str, options: List[str], default: str): + super().__init__() + self._step = step + self._question = question + self._options = options + self._default = default + + def on_mount(self) -> None: + self.push_screen( + PickScreen(self._step, self._question, self._options, self._default), + self.exit, + ) + + +def _ask(step: str, question: str, options: List[str], default: str) -> Optional[str]: + return _Ask(step, question, options, default).run() + + +def _detect_gpu() -> str: + if shutil.which("nvidia-smi"): + if subprocess.run(["nvidia-smi"], capture_output=True).returncode == 0: + return "nvidia" + if Path("/dev/dri").exists(): + return "mesa" + return "software" + + +def _docker_images() -> set: + r = subprocess.run( + ["docker", "images", "--format", "{{.Repository}}:{{.Tag}}"], + capture_output=True, text=True, + ) + return set(r.stdout.strip().splitlines()) + + +def _find_image(candidates: List[str]) -> Optional[str]: + available = _docker_images() + for img in candidates: + if img in available: + return img + return None + + +# --------------------------------------------------------------------------- +# Build project (colcon build --mixin release) +# --------------------------------------------------------------------------- + +def _task_build(screen: LogScreen) -> None: + try: + screen.write("[bold]Building project with colcon[/bold]\n") + + list_proc = subprocess.run( + ["bash", "-c", f"source {_JAZZY_DIR}/setup.bash && colcon list"], + capture_output=True, text=True, cwd=_PROJECT_DIR, + ) + total = max(len([l for l in list_proc.stdout.splitlines() if l.strip()]), 1) + screen.set_progress(0, f"0 / {total} packages done") + built = 0 + + proc = subprocess.Popen( + ["bash", "-c", + f"source {_JAZZY_DIR}/setup.bash && colcon build --mixin release"], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, cwd=_PROJECT_DIR, + ) + for line in proc.stdout: + s = line.rstrip() + if s: + screen.write(s) + if "Finished <<<" in line or "Failed <<<" in line: + built += 1 + screen.set_progress(built / total * 100, f"{built} / {total} packages done") + proc.wait() + + if proc.returncode != 0: + screen.write("\n[red]Build failed.[/red]") + screen.finish(False) + return + + screen.set_progress(100, "Build complete") + screen.write("\n[green]Build successful.[/green]") + screen.finish(True) + except Exception as exc: + screen.write(f"\n[red]Error:[/red] {exc}") + screen.finish(False) + + +class _BuildApp(App[bool]): + CSS = SCREEN_CSS + + def on_mount(self) -> None: + self.push_screen( + LogScreen("Building project before launch", _task_build, show_progress=True), + self.exit, + ) + + +# --------------------------------------------------------------------------- +# Local launch task +# --------------------------------------------------------------------------- + +def _task_run_local(screen: RunScreen, mode: str) -> None: + config = str(_CONFIG_PATH) + ros_cmd = f"ros2 launch iiwa_bringup iiwa.launch.py setting:={config}" + if mode == "webots": + ros_cmd += " simulate:=1" + + full_cmd = ( + f"source {_JAZZY_DIR}/setup.bash && " + f"source {_INSTALL_DIR}/setup.bash && " + f"{ros_cmd}" + ) + + label = "Webots simulator" if mode == "webots" else "Controller" + screen.write(f"[bold]Launching {label} (local)[/bold]") + screen.write(f"[dim]{ros_cmd}[/dim]\n") + + proc = subprocess.Popen( + ["bash", "-c", full_cmd], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, cwd=_PROJECT_DIR, + start_new_session=True, + ) + screen.set_proc(proc) + screen.set_kill_fn(lambda: os.killpg(os.getpgid(proc.pid), signal.SIGTERM)) + + for line in proc.stdout: + s = line.rstrip() + if s: + screen.write(s) + + proc.wait() + screen.finish(stopped=screen._stopped) + + +# --------------------------------------------------------------------------- +# Docker launch task +# --------------------------------------------------------------------------- + +def _task_run_docker(screen: RunScreen, image: str, mode: str, gpu: str) -> None: + container = _CONTAINER_WEBOTS if mode == "webots" else _CONTAINER_CONTROLLER + + ros_cmd = f"ros2 launch iiwa_bringup iiwa.launch.py setting:={_CONFIG_IN_CONTAINER}" + if mode == "webots": + ros_cmd += " simulate:=1" + + # Remove stale container with the same name + subprocess.run(["docker", "rm", "-f", container], capture_output=True) + + cmd = [ + "docker", "run", "--rm", + "--name", container, + "--network", "host", + "-e", "USER=root", + ] + + if mode == "webots": + subprocess.run(["xhost", "+local:docker"], capture_output=True) + cmd += [ + "-e", f"DISPLAY={os.environ.get('DISPLAY', ':0')}", + "-e", "QT_X11_NO_MITSHM=1", + "-v", "/tmp/.X11-unix:/tmp/.X11-unix:rw", + ] + if gpu == "nvidia": + cmd += [ + "--gpus", "all", + "-e", "NVIDIA_VISIBLE_DEVICES=all", + "-e", "NVIDIA_DRIVER_CAPABILITIES=graphics,utility,compute", + ] + elif gpu == "mesa": + cmd += ["--device", "/dev/dri"] + else: + cmd += [ + "-e", "LIBGL_ALWAYS_SOFTWARE=1", + "-e", "GALLIUM_DRIVER=llvmpipe", + ] + + if _CONFIG_PATH.exists(): + cmd += ["-v", f"{_CONFIG_PATH}:{_CONFIG_IN_CONTAINER}:ro"] + + cmd += [image, "bash", "-c", ros_cmd] + + _GPU_LABELS = { + "nvidia": "NVIDIA GPU", + "mesa": "Intel/AMD DRI (Mesa)", + "software": "Software rendering (llvmpipe)", + } + label = "Webots simulator" if mode == "webots" else "Controller" + screen.write(f"[bold]Launching {label} in Docker[/bold]") + screen.write(f"[dim]Image: {image}[/dim]") + if mode == "webots": + screen.write(f"[dim]GPU: {_GPU_LABELS.get(gpu, gpu)}[/dim]") + screen.write("") + + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, + ) + screen.set_proc(proc) + + for line in proc.stdout: + s = line.rstrip() + if s: + screen.write(s) + + proc.wait() + screen.finish(stopped=screen._stopped) + + +# --------------------------------------------------------------------------- +# RunApp wrapper +# --------------------------------------------------------------------------- + +class _RunApp(App[None]): + CSS = SCREEN_CSS + + def __init__(self, title: str, task: Callable): + super().__init__() + self._title = title + self._run_fn = task + + def on_mount(self) -> None: + self.push_screen(RunScreen(self._title, self._run_fn), lambda _: self.exit()) + + +# --------------------------------------------------------------------------- +# Local flow +# --------------------------------------------------------------------------- + +def _local_flow(args: argparse.Namespace) -> None: + mode_v = _ask( + "Run local", + "What do you want to launch?", + ["Controller", "Webots simulator"], + "Controller", + ) + if mode_v is None: + return + mode = "webots" if mode_v == "Webots simulator" else "controller" + + # Check ROS2 Jazzy + if not _JAZZY_DIR.is_dir(): + v = _ask( + "ROS2 not found", + "ROS2 Jazzy is not installed. Run local-setup now?", + ["Yes, run local-setup", "No, cancel"], + "Yes, run local-setup", + ) + if v and v.startswith("Yes"): + from cobot.commands.local_setup import run as _local_setup + _local_setup(args) + return + + # Check project built + if not (_INSTALL_DIR / "setup.bash").exists(): + v = _ask( + "Project not built", + "The project has not been built yet. Build it now?", + ["Yes, build now", "No, cancel"], + "Yes, build now", + ) + if v is None or v.startswith("No"): + return + ok = _BuildApp().run() + if not ok: + return + + label = "Webots simulator" if mode == "webots" else "Controller" + _RunApp(f"Running {label} — local", lambda s: _task_run_local(s, mode)).run() + + +# --------------------------------------------------------------------------- +# Docker flow +# --------------------------------------------------------------------------- + +def _docker_flow(args: argparse.Namespace) -> None: + if not shutil.which("docker"): + from rich.console import Console + Console().print("[red]Error:[/red] Docker is not installed or not on PATH.") + return + + mode_v = _ask( + "Run in Docker", + "What do you want to launch?", + ["Controller", "Webots simulator"], + "Controller", + ) + if mode_v is None: + return + mode = "webots" if mode_v == "Webots simulator" else "controller" + + candidates = _WEBOTS_IMAGES if mode == "webots" else _CONTROLLER_IMAGES + image = _find_image(candidates) + + if image is None: + what = "Webots" if mode == "webots" else "controller or Webots" + v = _ask( + "No image found", + f"No Docker image found for {what}. Run docker-setup now?", + ["Yes, run docker-setup", "No, cancel"], + "Yes, run docker-setup", + ) + if v and v.startswith("Yes"): + from cobot.commands.docker_setup import run as _docker_setup + _docker_setup(args) + return + + gpu = _detect_gpu() if mode == "webots" else "software" + + label = "Webots simulator" if mode == "webots" else "Controller" + _RunApp( + f"Running {label} — Docker", + lambda s: _task_run_docker(s, image, mode, gpu), + ).run() + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def register(subparsers: argparse._SubParsersAction) -> None: + p = subparsers.add_parser( + "run", + help="Launch the robot controller or Webots simulator", + ) + p.add_argument( + "mode", + nargs="?", + choices=["local", "docker"], + default=None, + help="local — native ROS2, docker — Docker container (default: ask)", + ) + p.set_defaults(func=run) + + +def run(args: argparse.Namespace) -> None: + mode = getattr(args, "mode", None) + + if mode == "local": + _local_flow(args) + elif mode == "docker": + _docker_flow(args) + else: + v = _ask( + "Run", + "How do you want to run the project?", + ["Local (native ROS2)", "Docker"], + "Local (native ROS2)", + ) + if v is None: + return + if v.startswith("Local"): + _local_flow(args) + else: + _docker_flow(args) diff --git a/cobot/tui.py b/cobot/tui.py index a621713..1b45076 100644 --- a/cobot/tui.py +++ b/cobot/tui.py @@ -59,6 +59,21 @@ LogScreen #hint { margin-top: 1; color: $text; } +RunScreen #log { + height: 1fr; + border: none; + padding: 0 1; + margin-top: 1; +} +RunScreen #loading { + height: 1; + margin-top: 1; +} +RunScreen #hint { + margin-top: 1; + color: $text-muted; + text-style: dim; +} """ @@ -149,6 +164,7 @@ class LogScreen(Screen[bool]): self._title = title self._run_fn = task self._finished = False + self._success = False self._show_progress = show_progress def compose(self) -> ComposeResult: @@ -188,6 +204,7 @@ class LogScreen(Screen[bool]): def _do_finish(self, success: bool) -> None: self._finished = True + self._success = success self.query_one("#loading", LoadingIndicator).display = False msg = ( "[green]Done![/green] Press [bold]Enter[/bold] to close." @@ -198,4 +215,77 @@ class LogScreen(Screen[bool]): def action_close(self) -> None: if self._finished: - self.dismiss(True) + self.dismiss(self._success) + + +class RunScreen(Screen[None]): + """Streams a long-running process. S/Enter/Escape stops or closes.""" + + BINDINGS = [ + Binding("s", "stop_close", "Stop", show=True, priority=True), + Binding("enter", "stop_close", "Close", show=False), + Binding("escape", "stop_close", "Close", show=False), + ] + + def __init__(self, title: str, task: Callable[[RunScreen], None]): + super().__init__() + self._title = title + self._run_fn = task + self._proc = None # set via set_proc() + self._kill_fn = None # optional custom kill callable + self._finished = False + self._stopped = False + + def compose(self) -> ComposeResult: + yield Static(self._title, id="step") + yield RichLog(id="log", highlight=True, markup=True, wrap=True) + yield LoadingIndicator(id="loading") + yield Static(" Press [bold]S[/bold] to stop the process", id="hint") + yield Footer() + + def on_mount(self) -> None: + self.query_one(RichLog).focus() + self.app.run_worker(lambda: self._run_fn(self), thread=True) + + def set_proc(self, proc) -> None: + """Register the running subprocess so Stop can terminate it.""" + self._proc = proc + + def set_kill_fn(self, fn: Callable) -> None: + """Override the default terminate() with a custom kill function.""" + self._kill_fn = fn + + def write(self, line: str) -> None: + self.app.call_from_thread(self._append, line) + + def _append(self, line: str) -> None: + self.query_one(RichLog).write(line) + + def finish(self, stopped: bool = False) -> None: + self.app.call_from_thread(self._do_finish, stopped) + + def _do_finish(self, stopped: bool) -> None: + self._finished = True + self.query_one("#loading", LoadingIndicator).display = False + if stopped: + msg = "[yellow]Process stopped.[/yellow] Press [bold]Enter[/bold] to close." + else: + msg = "[green]Process exited.[/green] Press [bold]Enter[/bold] to close." + self.query_one("#hint", Static).update(msg) + + def action_stop_close(self) -> None: + if self._finished: + self.dismiss(None) + return + self._stopped = True + if self._kill_fn is not None: + try: + self._kill_fn() + except Exception: + pass + elif self._proc is not None and self._proc.poll() is None: + try: + self._proc.terminate() + except Exception: + pass + self.write("\n[yellow]Stopping process...[/yellow]")