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:
Даниил Грабарь
2026-05-20 17:41:11 +03:00
parent bbfec17876
commit 5f16825cf7
6 changed files with 507 additions and 502 deletions
+139 -200
View File
@@ -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()