Feature/deploy misc

scripts/deploy/deploy_robots.py

@@ -52,7 +52,7 @@ def __init__(self):
     def _parse_arguments(self) -> argparse.Namespace:
         parser = ArgumentParserShowTargets(
             description="Deploy the Bit-Bots software on a robot. "
-            "This script provides 4 tasks: sync, configure, build, launch. "
+            "This script provides these tasks: sync, configure, build, launch. "
             "By default, it runs all tasks. You can select a subset of tasks by using the corresponding flags. "
             "For example, to only run the sync and build task, use the -sb."
         )
@@ -62,7 +62,7 @@ def _parse_arguments(self) -> argparse.Namespace:
             "targets",
             type=str,
             nargs="+",
-            help="The targets to deploy to. Multiple targets can be specified. 'ALL' can be used to target all known robots.",
+            help="The targets to deploy to. Multiple targets can be specified (SPACE separated). 'ALL' can be used to target all known robots. Examples: 'kalliope', '1', '<IP_ADDRESS>', 'ALL'",
         )
 
         parser.add_argument("--show-targets", action="store_true", help="Show all known targets and exit.")

scripts/deploy/known_targets.yaml

@@ -2,29 +2,37 @@
 "172.20.1.11":
   hostname: "kalliope"
   robot_name: "Kalliope"
+  robot_number: 1
 "172.20.1.12":
   hostname: "mickey"
   robot_name: "Mickey"
+  robot_number: 2
 "172.20.1.13":
   hostname: "pink"
   robot_name: "Pink"
+  robot_number: 3
 "172.20.1.14":
   hostname: "romeo"
   robot_name: "Romeo"
+  robot_number: 4
 
 # Field
 "10.0.6.1":
   hostname: "kalliope"
   robot_name: "Kalliope"
+  robot_number: 1
 "10.0.6.2":
   hostname: "mickey"
   robot_name: "Mickey"
+  robot_number: 2
 "10.0.6.3":
   hostname: "pink"
   robot_name: "Pink"
+  robot_number: 3
 "10.0.6.4":
   hostname: "romeo"
   robot_name: "Romeo"
+  robot_number: 4
 
 # HULKs
 #"10.1.24.125":

scripts/deploy/misc.py

@@ -94,7 +94,7 @@ def hide_output() -> bool | str:
 _known_targets_path: str = os.path.join(os.path.dirname(__file__), "known_targets.yaml")
 try:
     with open(_known_targets_path) as f:
-        KNOWN_TARGETS: dict[str, dict[str, str]] = yaml.safe_load(f)
+        KNOWN_TARGETS: dict[str, dict[str, str | int]] = yaml.safe_load(f)
 except FileNotFoundError:
     print_error(f"Could not find known_targets.yaml in {_known_targets_path}")
     sys.exit(1)
@@ -104,18 +104,19 @@ def print_known_targets() -> None:
     table = Table()
     table.add_column("Hostname")
     table.add_column("Robot name")
+    table.add_column("Robot number")
     table.add_column("IP address")
 
     known_targets = get_known_targets()
-    table.add_row("ALL", "", "")
+    table.add_row("ALL", "", "", "")
     for ip, values in known_targets.items():
-        table.add_row(values.get("hostname", ""), values.get("robot_name", ""), ip)
+        table.add_row(values.get("hostname", ""), values.get("robot_name", ""), str(values.get("robot_number", "")), ip)
     print_info("You can enter the following values as targets:")
     CONSOLE.print(table)
     sys.exit(0)
 
 
-def get_known_targets() -> dict[str, dict[str, str]]:
+def get_known_targets() -> dict[str, dict[str, str | int]]:
     """
     Returns the known targets.
 
@@ -171,37 +172,70 @@ def append_target_ip(target_ips: list[str], ip: str) -> None:
             else:
                 print_debug(f"Robot name '{known_robot_name}' does not match '{identifier}'.")
 
+        # Is the identifier a known robot number?
+        known_robot_number = values.get("robot_number", None)
+        if known_robot_number is not None:
+            print_debug(f"Comparing '{identifier}' with robot number '{known_robot_number}'.")
+            if str(known_robot_number).strip() == identifier.strip():
+                print_debug(f"Identified robot number '{known_robot_number}' for '{identifier}'. Using its IP {ip}.")
+                append_target_ip(target_ips, ip)
+            else:
+                print_debug(f"Robot number '{known_robot_number}' does not match '{identifier}'.")
+
     return target_ips
 
 
 def _parse_targets(targets: list[str]) -> list[str]:
     """
     Parse target input into usable target IP addresses.
 
-    :param targets: Targets as a list of strings of either hostnames, robot names or IPs.
+    :param targets: Targets as a list of strings of either hostnames, robot names, robot numbers or IPs.
     :return: List of target IP addresses.
     """
     return [target_ip for _, target_ips in _parse_target_candidates(targets) for target_ip in target_ips]
 
 
+def _identify_ips_or_exit(target: str) -> list[str]:
+    """
+    Identifies IP addresses from an identifier or exits the program if no IP address could be identified
+
+    :param target: The identifier to identify the target from.
+    :return: IP addresses of the identified targets.
+    """
+    try:
+        target_ips = _identify_ips(target)
+    except ValueError:
+        print_error(f"Could not determine IP address from input: '{target}'")
+        sys.exit(1)
+
+    if not target_ips:
+        print_error(f"Could not determine IP address from input: '{target}'")
+        sys.exit(1)
+
+    return target_ips
+
+
 def _parse_target_candidates(targets: list[str]) -> list[tuple[str, list[str]]]:
     """
     Parse target input into candidate target IP addresses grouped by user input.
 
-    :param targets: Targets as a list of strings of either hostnames, robot names or IPs.
+    :param targets: Targets as a list of strings of either hostnames, robot names, robot numbers or IPs.
     :return: List of target names and their candidate IP addresses.
     """
     target_candidates: list[tuple[str, list[str]]] = []
+    seen_target_ips: set[str] = set()
     for target in targets:
-        try:
-            target_ips_for_identifier = _identify_ips(target)
-        except ValueError:
-            print_error(f"Could not determine IP address from input: '{target}'")
-            sys.exit(1)
-        if not target_ips_for_identifier:
-            print_error(f"Could not determine IP address from input: '{target}'")
-            sys.exit(1)
-        target_candidates.append((target, target_ips_for_identifier))
+        target_ips = []
+        for target_ip in _identify_ips_or_exit(target):
+            if target_ip in seen_target_ips:
+                continue
+            seen_target_ips.add(target_ip)
+            target_ips.append(target_ip)
+
+        if target_ips:
+            target_candidates.append((target, target_ips))
+        else:
+            print_debug(f"Target '{target}' only resolved to already requested IPs. Skipping it.")
     return target_candidates
 
 
@@ -246,6 +280,7 @@ def _try_connect(target_ip: str) -> tuple[str, Connection | None, str | None]:
             return target_ip, None, _concat_exception_args(e)
 
     open_connections: list[Connection] = []
+    connected_target_ips: set[str] = set()
     failed_targets: list[tuple[str, list[tuple[str, str]]]] = []
 
     for target, target_ips in target_candidates:
@@ -273,7 +308,12 @@ def _try_connect(target_ip: str) -> tuple[str, Connection | None, str | None]:
             selected_target_ip = successful_target_ips[0]
             selected_connection = results_by_ip[selected_target_ip][0]
             assert selected_connection is not None
-            open_connections.append(selected_connection)
+            if selected_target_ip in connected_target_ips:
+                print_debug(f"Already connected to {selected_target_ip}. Closing duplicate connection.")
+                selected_connection.close()
+            else:
+                open_connections.append(selected_connection)
+                connected_target_ips.add(selected_target_ip)
             for target_ip in successful_target_ips[1:]:
                 ignored_connection = results_by_ip[target_ip][0]
                 assert ignored_connection is not None

scripts/deploy/tasks/launch.py

@@ -1,9 +1,19 @@
-import re
-
-from deploy.misc import get_connections_from_succeeded, hide_output, print_debug, print_error, print_success
+from deploy.misc import (
+    CONSOLE,
+    get_connections_from_succeeded,
+    hide_output,
+    print_debug,
+    print_error,
+    print_success,
+    print_warning,
+)
 from deploy.tasks.abstract_task import AbstractTask
-from fabric import Group, GroupResult, Result
+from fabric import Group, GroupResult, Result, ThreadingGroup
 from fabric.exceptions import GroupException
+from rich import box
+from rich.console import Group as RichGroup
+from rich.prompt import Confirm
+from rich.table import Table
 
 
 class Launch(AbstractTask):
@@ -15,6 +25,7 @@ def __init__(self, remote_workspace: str, tmux_session_name: str) -> None:
         :param tmux_session_name: Name of the fresh tmux session to launch the teamplayer in
         """
         super().__init__()
+        self._show_status = False
 
         self._remote_workspace = remote_workspace
         self._tmux_session_name = tmux_session_name
@@ -31,20 +42,57 @@ def _run(self, connections: Group) -> GroupResult:
         :return: The results of the task.
         """
         # Check if ROS 2 nodes are already running
-        node_running_results = self._check_nodes_already_running(connections)
-        if not node_running_results.succeeded:
+        with CONSOLE.status("[bold blue] Checking if ROS 2 nodes are already running", spinner="point"):
+            node_running_results = self._check_nodes_already_running(connections)
+        if node_running_results.failed:
+            self._show_running_entries(node_running_results.failed, "Running ROS 2 nodes", "Nodes")
+            if not Confirm.ask("Kill running ROS 2 nodes on these hosts?", console=CONSOLE, default=False):
+                return node_running_results
+            kill_nodes_results = self._kill_running_ros_nodes(
+                self._connections_from_results(node_running_results.failed)
+            )
+            if kill_nodes_results.failed:
+                return kill_nodes_results
+            with CONSOLE.status("[bold blue] Checking if ROS 2 nodes stopped", spinner="point"):
+                node_running_results = self._check_nodes_already_running(connections)
+        if node_running_results.failed:
             return node_running_results
         # Some hosts have no ROS 2 nodes already running, continuing
 
         # Check if tmux session is already running
-        tmux_session_running_results = self._check_tmux_session_already_running(
-            get_connections_from_succeeded(node_running_results)
-        )
-        if not tmux_session_running_results.succeeded:
+        with CONSOLE.status(
+            f"[bold blue] Checking if tmux session '{self._tmux_session_name}' is already running",
+            spinner="point",
+        ):
+            tmux_session_running_results = self._check_tmux_session_already_running(
+                get_connections_from_succeeded(node_running_results)
+            )
+        if tmux_session_running_results.failed:
+            self._show_running_entries(tmux_session_running_results.failed, "Running tmux sessions", "Sessions")
+            if not Confirm.ask(
+                f"Kill tmux session '{self._tmux_session_name}' on these hosts?",
+                console=CONSOLE,
+                default=False,
+            ):
+                return tmux_session_running_results
+            kill_tmux_results = self._kill_running_tmux_session(
+                self._connections_from_results(tmux_session_running_results.failed)
+            )
+            if kill_tmux_results.failed:
+                return kill_tmux_results
+            with CONSOLE.status(
+                f"[bold blue] Checking if tmux session '{self._tmux_session_name}' stopped",
+                spinner="point",
+            ):
+                tmux_session_running_results = self._check_tmux_session_already_running(
+                    get_connections_from_succeeded(node_running_results)
+                )
+        if tmux_session_running_results.failed:
             return tmux_session_running_results
 
         # Some hosts are ready to launch teamplayer
-        launch_results = self._launch_teamplayer(get_connections_from_succeeded(tmux_session_running_results))
+        with CONSOLE.status("[bold blue] Launching teamplayer", spinner="point"):
+            launch_results = self._launch_teamplayer(get_connections_from_succeeded(tmux_session_running_results))
         return launch_results
 
     def _check_nodes_already_running(self, connections: Group) -> GroupResult:
@@ -55,7 +103,7 @@ def _check_nodes_already_running(self, connections: Group) -> GroupResult:
         :return: Results, with success if ROS 2 nodes are not already running
         """
         print_debug("Checking if ROS 2 nodes are already running")
-        cmd = f"cd {self._remote_workspace} && pixi run --environment robot ros2 node list -c"
+        cmd = f"cd {self._remote_workspace} && pixi run --environment robot ros2 node list"
 
         print_debug(f"Calling '{cmd}'")
         try:
@@ -68,15 +116,12 @@ def _check_nodes_already_running(self, connections: Group) -> GroupResult:
             else:
                 node_list_results = e.result
 
-        # Check if output was ^0$ (zero for no nodes running) for all remote machines
+        # Check if the node list output is empty for all remote machines
         # Create new group result with success if no nodes are running
         no_nodes_already_running_results = GroupResult()
         nodes_already_running: bool
         for connection, result in node_list_results.items():
-            if re.match(r"^0$", result.stdout):  # No nodes already running
-                nodes_already_running = False
-            else:  # Nodes already running
-                nodes_already_running = True
+            nodes_already_running = bool(result.stdout.strip())
 
             # Create new result and pass most of the data from the original result.
             # Exited is the exit code.
@@ -97,6 +142,33 @@ def _check_nodes_already_running(self, connections: Group) -> GroupResult:
             print_error(f"ROS 2 nodes are already running on {self._failed_hosts(no_nodes_already_running_results)}")
         return no_nodes_already_running_results
 
+    def _kill_running_ros_nodes(self, connections: Group) -> GroupResult:
+        print_debug("Killing running ROS 2 nodes")
+        script = (
+            'nodes="$(ros2 node list 2>/dev/null || true)"; '
+            "while IFS= read -r node; do "
+            '[ -n "$node" ] && ros2 lifecycle set "$node" shutdown >/dev/null 2>&1 || true; '
+            'done <<< "$nodes"; '
+            'pattern="ros""2|launch_""ros|component_""container|robot_state_""publisher"; '
+            'pgrep -u "$USER" -f "$pattern" | while read -r pid; do '
+            '[ "$pid" != "$$" ] && kill -INT "$pid" 2>/dev/null || true; '
+            "done; "
+            "sleep 2; "
+            'pgrep -u "$USER" -f "$pattern" | while read -r pid; do '
+            '[ "$pid" != "$$" ] && kill -TERM "$pid" 2>/dev/null || true; '
+            "done"
+        )
+        cmd = f"cd {self._remote_workspace} && pixi run --environment robot bash -lc '{script}'"
+
+        print_debug(f"Calling '{cmd}'")
+        try:
+            results = connections.run(cmd, hide=hide_output())
+            print_debug(f"Calling '{cmd}' succeeded on: {self._succeeded_hosts(results)}")
+        except GroupException as e:
+            print_error(f"Killing running ROS 2 nodes failed on: {self._failed_hosts(e.result)}")
+            return e.result
+        return results
+
     def _check_tmux_session_already_running(self, connections: Group) -> GroupResult:
         print_debug(f"Checking if tmux session with name '{self._tmux_session_name}' is already running")
 
@@ -149,6 +221,34 @@ def _check_tmux_session_already_running(self, connections: Group) -> GroupResult
             )
         return tmux_session_already_running_results
 
+    def _kill_running_tmux_session(self, connections: Group) -> GroupResult:
+        print_debug(f"Killing tmux session with name '{self._tmux_session_name}'")
+        cmd = f"tmux kill-session -t {self._tmux_session_name}"
+
+        print_debug(f"Calling '{cmd}'")
+        try:
+            results = connections.run(cmd, hide=hide_output())
+            print_debug(f"Calling '{cmd}' succeeded on: {self._succeeded_hosts(results)}")
+        except GroupException as e:
+            print_error(
+                f"Killing tmux session called {self._tmux_session_name} failed on: {self._failed_hosts(e.result)}"
+            )
+            return e.result
+        return results
+
+    def _show_running_entries(self, results: GroupResult, title: str, value_column: str) -> None:
+        table = Table(style="yellow", box=box.SQUARE)
+        table.add_column("Host")
+        table.add_column(value_column)
+
+        for connection, result in results.items():
+            table.add_row(connection.original_host, result.stdout.strip())
+
+        print_warning(RichGroup(title, table))
+
+    def _connections_from_results(self, results: GroupResult) -> ThreadingGroup:
+        return ThreadingGroup.from_connections(results.keys())
+
     def _launch_teamplayer(self, connections: Group) -> GroupResult:
         print_debug("Launching teamplayer")
         # Create tmux session