Refactor Docker and Local Setup Commands
- Enhanced the docker_setup.py to streamline image building and pulling processes with improved logging and error handling. - Introduced a new LogScreen for better user feedback during long-running tasks. - Updated local_setup.py to utilize logging for installation steps and improved error handling. - Refactored robot_setup.py to simplify configuration saving and user interaction. - Added a new LogScreen class in tui.py for consistent logging across different setup processes. - Modified install.sh to adjust the installation directory for better organization.
This commit is contained in:
+139
-200
@@ -8,15 +8,11 @@ import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
from rich.console import Console
|
||||
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
|
||||
from textual.app import App
|
||||
|
||||
from cobot.tui import SCREEN_CSS, InputScreen, PickScreen
|
||||
|
||||
_console = Console()
|
||||
from cobot.tui import SCREEN_CSS, InputScreen, LogScreen, PickScreen
|
||||
|
||||
_PROJECT_DIR = Path(__file__).parent.parent.parent
|
||||
_DOCKER_DIR = _PROJECT_DIR / "docker"
|
||||
@@ -33,22 +29,128 @@ _IMAGE_PARENT: dict[str, str | None] = {
|
||||
"ros-iiwa7-webots": "ros-base",
|
||||
}
|
||||
|
||||
# Images that COPY from src/ — need project root as build context.
|
||||
# Others use their Dockerfile's own directory.
|
||||
_NEEDS_PROJECT_CTX = {"ros-iiwa7", "ros-iiwa7-webots"}
|
||||
|
||||
Write = Callable[[str], None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Config:
|
||||
ros_version: str
|
||||
variant: str # "controller" | "webots"
|
||||
source: str # "build" | "pull"
|
||||
build_type: str # "release" | "dev"
|
||||
variant: str
|
||||
source: str
|
||||
build_type: str
|
||||
image_prefix: str
|
||||
hub_repo: str
|
||||
|
||||
|
||||
class _Wizard(App[Optional[_Config]]):
|
||||
def _stream(cmd: List[str], write: Write, env=None) -> bool:
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||||
text=True, env=env or os.environ,
|
||||
)
|
||||
for line in proc.stdout:
|
||||
s = line.rstrip()
|
||||
if s:
|
||||
write(s)
|
||||
proc.wait()
|
||||
return proc.returncode == 0
|
||||
|
||||
|
||||
def _build_image(
|
||||
name: str, tag: str, dockerfile: Path, ctx: Path,
|
||||
write: Write,
|
||||
parent_tag: Optional[str] = None,
|
||||
build_type: str = "release",
|
||||
) -> bool:
|
||||
write(f"[cyan][*][/cyan] Building [bold]{name}[/bold]...")
|
||||
env = {**os.environ, "DOCKER_BUILDKIT": "0"}
|
||||
cmd = [
|
||||
"docker", "build", "-t", tag, "-f", str(dockerfile),
|
||||
"--build-arg", f"BUILD_TYPE={build_type}",
|
||||
]
|
||||
if parent_tag:
|
||||
cmd += ["--build-arg", f"IMAGE={parent_tag}"]
|
||||
cmd.append(str(ctx))
|
||||
ok = _stream(cmd, write, env)
|
||||
if ok:
|
||||
write(f"[green][ok][/green] {name}")
|
||||
else:
|
||||
write(f"[red]Build failed:[/red] {tag}")
|
||||
return ok
|
||||
|
||||
|
||||
def _pull_image(name: str, tag: str, write: Write) -> bool:
|
||||
write(f"[cyan][*][/cyan] Pulling [bold]{name}[/bold] ({tag})...")
|
||||
ok = _stream(["docker", "pull", tag], write)
|
||||
if ok:
|
||||
write(f"[green][ok][/green] {name}")
|
||||
else:
|
||||
write(f"[red]Pull failed:[/red] {tag}")
|
||||
return ok
|
||||
|
||||
|
||||
def _task_execute(screen: LogScreen, cfg: _Config) -> None:
|
||||
try:
|
||||
chain = _WEBOTS_CHAIN if cfg.variant == "webots" else _CONTROLLER_CHAIN
|
||||
|
||||
if cfg.source == "build":
|
||||
screen.write(
|
||||
f"[bold]Building {len(chain)} image(s) — "
|
||||
f"ROS {cfg.ros_version} — {cfg.build_type}[/bold]\n"
|
||||
)
|
||||
for name in chain:
|
||||
tag = f"{cfg.image_prefix}:{name}-{cfg.ros_version}"
|
||||
dockerfile = _DOCKER_DIR / cfg.ros_version / name / "Dockerfile"
|
||||
if not dockerfile.exists():
|
||||
screen.write(f"[red]Dockerfile not found:[/red] {dockerfile}")
|
||||
screen.finish(False)
|
||||
return
|
||||
ctx = _PROJECT_DIR if name in _NEEDS_PROJECT_CTX else dockerfile.parent
|
||||
parent_name = _IMAGE_PARENT.get(name)
|
||||
parent_tag = (
|
||||
f"{cfg.image_prefix}:{parent_name}-{cfg.ros_version}"
|
||||
if parent_name else None
|
||||
)
|
||||
if not _build_image(name, tag, dockerfile, ctx, screen.write, parent_tag, cfg.build_type):
|
||||
screen.finish(False)
|
||||
return
|
||||
screen.write(
|
||||
f"\n[green]Done.[/green] "
|
||||
f"Images tagged [bold]{cfg.image_prefix}:<name>-{cfg.ros_version}[/bold]."
|
||||
)
|
||||
|
||||
else:
|
||||
short = "webots" if cfg.variant == "webots" else "iiwa"
|
||||
suffix = "-dev" if cfg.build_type == "dev" else ""
|
||||
full_ref = f"{cfg.hub_repo}:{short}-{cfg.ros_version}{suffix}"
|
||||
screen.write(
|
||||
f"[bold]Pulling from {cfg.hub_repo} — "
|
||||
f"ROS {cfg.ros_version} — {cfg.build_type}[/bold]\n"
|
||||
)
|
||||
if not _pull_image(short, full_ref, screen.write):
|
||||
screen.finish(False)
|
||||
return
|
||||
screen.write(f"\n[green]Done.[/green] Image ready: [bold]{full_ref}[/bold].")
|
||||
|
||||
screen.finish(True)
|
||||
|
||||
except Exception as exc:
|
||||
screen.write(f"\n[red]Error:[/red] {exc}")
|
||||
screen.finish(False)
|
||||
|
||||
|
||||
def _discover_versions() -> List[str]:
|
||||
if not _DOCKER_DIR.exists():
|
||||
return ["jazzy"]
|
||||
dirs = sorted(d.name for d in _DOCKER_DIR.iterdir() if d.is_dir())
|
||||
if "jazzy" in dirs:
|
||||
dirs = ["jazzy"] + [d for d in dirs if d != "jazzy"]
|
||||
return dirs or ["jazzy"]
|
||||
|
||||
|
||||
class _Wizard(App[None]):
|
||||
CSS = SCREEN_CSS
|
||||
|
||||
def __init__(self, versions: List[str], default_version: str = "jazzy"):
|
||||
@@ -60,8 +162,6 @@ class _Wizard(App[Optional[_Config]]):
|
||||
def on_mount(self) -> None:
|
||||
self._ask_version()
|
||||
|
||||
# ── step helpers ──────────────────────────────────────────────────────────
|
||||
|
||||
def _ask_version(self) -> None:
|
||||
self.push_screen(
|
||||
PickScreen("Step 1 of 5", "Select ROS version:", self.versions, self.default_version),
|
||||
@@ -70,17 +170,21 @@ class _Wizard(App[Optional[_Config]]):
|
||||
|
||||
def _got_version(self, v: Optional[str]) -> None:
|
||||
if v is None:
|
||||
self.exit(None)
|
||||
self.exit()
|
||||
return
|
||||
self._state["ros_version"] = v
|
||||
self.push_screen(
|
||||
PickScreen("Step 2 of 5", "Source:", ["Pull from Docker Hub", "Build locally"], "Pull from Docker Hub"),
|
||||
PickScreen(
|
||||
"Step 2 of 5", "Source:",
|
||||
["Pull from Docker Hub", "Build locally"],
|
||||
"Pull from Docker Hub",
|
||||
),
|
||||
self._got_source,
|
||||
)
|
||||
|
||||
def _got_source(self, v: Optional[str]) -> None:
|
||||
if v is None:
|
||||
self.exit(None)
|
||||
self.exit()
|
||||
return
|
||||
self._state["source"] = "build" if v == "Build locally" else "pull"
|
||||
self.push_screen(
|
||||
@@ -98,7 +202,7 @@ class _Wizard(App[Optional[_Config]]):
|
||||
|
||||
def _got_variant(self, v: Optional[str]) -> None:
|
||||
if v is None:
|
||||
self.exit(None)
|
||||
self.exit()
|
||||
return
|
||||
self._state["variant"] = "webots" if v.startswith("Controller with Webots") else "controller"
|
||||
self.push_screen(
|
||||
@@ -108,7 +212,7 @@ class _Wizard(App[Optional[_Config]]):
|
||||
|
||||
def _got_build_type(self, v: Optional[str]) -> None:
|
||||
if v is None:
|
||||
self.exit(None)
|
||||
self.exit()
|
||||
return
|
||||
self._state["build_type"] = v or "release"
|
||||
if self._state["source"] == "pull":
|
||||
@@ -124,215 +228,50 @@ class _Wizard(App[Optional[_Config]]):
|
||||
|
||||
def _got_hub_repo(self, v: Optional[str]) -> None:
|
||||
if v is None:
|
||||
self.exit(None)
|
||||
self.exit()
|
||||
return
|
||||
self._state["hub_repo"] = v
|
||||
self._finish()
|
||||
|
||||
def _got_prefix(self, v: Optional[str]) -> None:
|
||||
if v is None:
|
||||
self.exit(None)
|
||||
self.exit()
|
||||
return
|
||||
self._state["image_prefix"] = v
|
||||
self._finish()
|
||||
|
||||
def _finish(self) -> None:
|
||||
s = self._state
|
||||
self.exit(_Config(
|
||||
cfg = _Config(
|
||||
ros_version=s["ros_version"],
|
||||
variant=s["variant"],
|
||||
source=s["source"],
|
||||
build_type=s["build_type"],
|
||||
image_prefix=s.get("image_prefix", _DEFAULT_PREFIX),
|
||||
hub_repo=s.get("hub_repo", _DEFAULT_HUB_REPO),
|
||||
))
|
||||
|
||||
|
||||
def _discover_versions() -> List[str]:
|
||||
if not _DOCKER_DIR.exists():
|
||||
return ["jazzy"]
|
||||
dirs = sorted(d.name for d in _DOCKER_DIR.iterdir() if d.is_dir())
|
||||
if "jazzy" in dirs:
|
||||
dirs = ["jazzy"] + [d for d in dirs if d != "jazzy"]
|
||||
return dirs or ["jazzy"]
|
||||
|
||||
|
||||
def _build_image(
|
||||
name: str, tag: str, dockerfile: Path, ctx: Path,
|
||||
parent_tag: Optional[str] = None,
|
||||
build_type: str = "release",
|
||||
) -> bool:
|
||||
env = {**os.environ, "DOCKER_BUILDKIT": "0"}
|
||||
cmd = [
|
||||
"docker", "build",
|
||||
"-t", tag,
|
||||
"-f", str(dockerfile),
|
||||
"--build-arg", f"BUILD_TYPE={build_type}",
|
||||
]
|
||||
if parent_tag:
|
||||
cmd += ["--build-arg", f"IMAGE={parent_tag}"]
|
||||
cmd.append(str(ctx))
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||||
text=True, env=env,
|
||||
)
|
||||
|
||||
captured: List[str] = []
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn(f" [bold cyan]{name:<24}[/bold cyan]"),
|
||||
BarColumn(bar_width=32),
|
||||
TaskProgressColumn(),
|
||||
console=_console,
|
||||
transient=False,
|
||||
) as prog:
|
||||
task = prog.add_task("", total=100)
|
||||
total = 1
|
||||
for line in proc.stdout:
|
||||
captured.append(line)
|
||||
m = re.match(r"Step (\d+)/(\d+) :", line)
|
||||
if m:
|
||||
step, total = int(m.group(1)), int(m.group(2))
|
||||
prog.update(task, completed=step / total * 100)
|
||||
prog.update(task, completed=100)
|
||||
|
||||
proc.wait()
|
||||
if proc.returncode != 0:
|
||||
_console.print(f"\n[red]Build failed:[/red] {tag}\n")
|
||||
_console.print("".join(captured), highlight=False)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _parse_docker_size(s: str) -> float:
|
||||
s = s.strip()
|
||||
for suffix, mult in [("GB", 1e9), ("MB", 1e6), ("kB", 1e3), ("B", 1.0)]:
|
||||
if s.endswith(suffix):
|
||||
try:
|
||||
return float(s[: -len(suffix)]) * mult
|
||||
except ValueError:
|
||||
return 0.0
|
||||
try:
|
||||
return float(s)
|
||||
except ValueError:
|
||||
return 0.0
|
||||
|
||||
|
||||
_RE_PULLING = re.compile(r"^([a-f0-9]+): Pulling fs layer")
|
||||
_RE_DONE = re.compile(r"^([a-f0-9]+): (?:Pull complete|Already exists|Layer already exists)")
|
||||
_RE_DL = re.compile(r"^([a-f0-9]+): Downloading(?:\s+\[.*?\])?\s+([\d.]+\s*\w+)/([\d.]+\s*\w+)")
|
||||
|
||||
|
||||
def _pull_image(name: str, tag: str) -> bool:
|
||||
proc = subprocess.Popen(
|
||||
["docker", "pull", tag],
|
||||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
)
|
||||
|
||||
captured: List[str] = []
|
||||
layers_pulling: set[str] = set()
|
||||
layers_done: set[str] = set()
|
||||
layers_total: dict[str, float] = {}
|
||||
layers_current: dict[str, float] = {}
|
||||
has_bytes = False
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn(f" [bold cyan]{name:<24}[/bold cyan]"),
|
||||
BarColumn(bar_width=32),
|
||||
TaskProgressColumn(),
|
||||
console=_console,
|
||||
transient=False,
|
||||
) as prog:
|
||||
task = prog.add_task("", total=100)
|
||||
|
||||
for line in proc.stdout:
|
||||
captured.append(line)
|
||||
line = line.strip()
|
||||
|
||||
if m := _RE_PULLING.match(line):
|
||||
layers_pulling.add(m.group(1))
|
||||
|
||||
elif m := _RE_DL.match(line):
|
||||
lid, cur, tot = m.group(1), _parse_docker_size(m.group(2)), _parse_docker_size(m.group(3))
|
||||
if tot > 0:
|
||||
has_bytes = True
|
||||
layers_current[lid] = cur
|
||||
layers_total[lid] = tot
|
||||
|
||||
elif m := _RE_DONE.match(line):
|
||||
lid = m.group(1)
|
||||
layers_done.add(lid)
|
||||
if lid in layers_total:
|
||||
layers_current[lid] = layers_total[lid]
|
||||
|
||||
if has_bytes and layers_total:
|
||||
total_b = sum(layers_total.values())
|
||||
done_b = sum(layers_current.get(lid, 0.0) for lid in layers_total)
|
||||
prog.update(task, completed=done_b / total_b * 100)
|
||||
elif layers_pulling:
|
||||
prog.update(task, completed=len(layers_done) / len(layers_pulling) * 100)
|
||||
|
||||
prog.update(task, completed=100)
|
||||
|
||||
proc.wait()
|
||||
if proc.returncode != 0:
|
||||
_console.print(f"\n[red]Pull failed:[/red] {tag}\n")
|
||||
_console.print("".join(captured), highlight=False)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _execute(cfg: _Config) -> None:
|
||||
chain = _WEBOTS_CHAIN if cfg.variant == "webots" else _CONTROLLER_CHAIN
|
||||
_console.print()
|
||||
|
||||
if cfg.source == "build":
|
||||
_console.print(
|
||||
f"[bold]Building {len(chain)} image(s) • ROS {cfg.ros_version} • {cfg.build_type}[/bold]\n"
|
||||
)
|
||||
for name in chain:
|
||||
tag = f"{cfg.image_prefix}:{name}-{cfg.ros_version}"
|
||||
dockerfile = _DOCKER_DIR / cfg.ros_version / name / "Dockerfile"
|
||||
if not dockerfile.exists():
|
||||
_console.print(f"[red]Dockerfile not found:[/red] {dockerfile}")
|
||||
sys.exit(1)
|
||||
ctx = _PROJECT_DIR if name in _NEEDS_PROJECT_CTX else dockerfile.parent
|
||||
parent_name = _IMAGE_PARENT.get(name)
|
||||
parent_tag = f"{cfg.image_prefix}:{parent_name}-{cfg.ros_version}" if parent_name else None
|
||||
if not _build_image(name, tag, dockerfile, ctx, parent_tag, cfg.build_type):
|
||||
sys.exit(1)
|
||||
_console.print(f"\n[green]Done.[/green] Images: [bold]{cfg.image_prefix}:{{name}}-{cfg.ros_version}[/bold].")
|
||||
|
||||
else:
|
||||
short = "webots" if cfg.variant == "webots" else "iiwa"
|
||||
suffix = "-dev" if cfg.build_type == "dev" else ""
|
||||
image_tag = f"{short}-{cfg.ros_version}{suffix}"
|
||||
full_ref = f"{cfg.hub_repo}:{image_tag}"
|
||||
_console.print(f"[bold]Pulling from {cfg.hub_repo} • ROS {cfg.ros_version} • {cfg.build_type}[/bold]\n")
|
||||
if not _pull_image(short, full_ref):
|
||||
sys.exit(1)
|
||||
_console.print(f"\n[green]Done.[/green] Image ready: [bold]{full_ref}[/bold].")
|
||||
title = (
|
||||
f"Building Docker images — ROS {cfg.ros_version}"
|
||||
if cfg.source == "build"
|
||||
else f"Pulling Docker image — ROS {cfg.ros_version}"
|
||||
)
|
||||
self.push_screen(
|
||||
LogScreen(title, lambda screen: _task_execute(screen, cfg)),
|
||||
lambda _: self.exit(),
|
||||
)
|
||||
|
||||
|
||||
def register(subparsers: argparse._SubParsersAction) -> None:
|
||||
p = subparsers.add_parser("docker-setup", help="Set up Docker images for KUKA iiwa7")
|
||||
p = subparsers.add_parser("docker-setup", help="Build or pull Docker images for KUKA iiwa7")
|
||||
p.set_defaults(func=run)
|
||||
|
||||
|
||||
def run(args: argparse.Namespace) -> None:
|
||||
if not shutil.which("docker"):
|
||||
_console.print("[red]Error:[/red] Docker is not installed or not on PATH.")
|
||||
from rich.console import Console
|
||||
Console().print("[red]Error:[/red] Docker is not installed or not on PATH.")
|
||||
sys.exit(1)
|
||||
|
||||
versions = _discover_versions()
|
||||
default = "jazzy" if "jazzy" in versions else versions[0]
|
||||
cfg = _Wizard(versions=versions, default_version=default).run()
|
||||
|
||||
if cfg is None:
|
||||
return
|
||||
|
||||
_execute(cfg)
|
||||
_Wizard(versions=versions, default_version=default).run()
|
||||
|
||||
Reference in New Issue
Block a user