ria-toolkit-oss/src/ria_toolkit_oss/app/cli.py

"""Unified ``ria-app`` CLI.

Subcommands:

- ``ria-app pull <app>[:tag]`` — pull a RIA app image from the configured registry.
- ``ria-app run  <app>[:tag]`` — pull (if needed) and run, auto-configuring
  GPU/USB/network flags from image labels set by CI.
- ``ria-app list`` — list locally cached RIA app images.
- ``ria-app stop <app>`` — stop a running app container.
- ``ria-app logs <app>`` — tail logs of a running app container.
- ``ria-app configure`` — set default registry/namespace.

Image references resolve as::

    my-classifier            -> {registry}/{namespace}/my-classifier:latest
    group/my-classifier      -> {registry}/group/my-classifier:latest
    host/group/app:tag       -> host/group/app:tag   (fully-qualified passthrough)
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys

from . import config as _config

_LABEL_PROFILE = "ria.profile"
_LABEL_HARDWARE = "ria.hardware"
_LABEL_APP = "ria.app"


def _engine(cfg: _config.AppConfig, sudo_override: bool = False) -> list[str]:
    for exe in ("docker", "podman"):
        if shutil.which(exe):
            use_sudo = sudo_override or cfg.sudo
            return (["sudo", exe] if use_sudo else [exe])
    print("error: neither 'docker' nor 'podman' found on PATH", file=sys.stderr)
    sys.exit(2)


def _resolve_ref(app: str, cfg: _config.AppConfig) -> str:
    ref = app if ":" in app.split("/")[-1] else f"{app}:latest"
    slashes = ref.count("/")
    if slashes >= 2:
        return ref
    if slashes == 1:
        return f"{cfg.registry}/{ref}" if cfg.registry else ref
    if not cfg.registry or not cfg.namespace:
        print(
            "error: app is not fully qualified and no default registry/namespace configured. "
            "Run `ria-app configure` or pass a full image reference (registry/namespace/app:tag).",
            file=sys.stderr,
        )
        sys.exit(2)
    return f"{cfg.registry}/{cfg.namespace}/{ref}"


def _container_name(ref: str) -> str:
    name = ref.rsplit("/", 1)[-1].split(":", 1)[0]
    return f"ria-app-{name}"


def _inspect_labels(engine: list[str], ref: str) -> dict:
    try:
        out = subprocess.check_output(
            [*engine, "image", "inspect", "--format", "{{json .Config.Labels}}", ref],
            stderr=subprocess.DEVNULL,
        )
    except subprocess.CalledProcessError:
        return {}
    try:
        return json.loads(out.decode().strip()) or {}
    except json.JSONDecodeError:
        return {}


def _gpu_available() -> bool:
    if os.path.exists("/dev/nvidia0"):
        return True
    return shutil.which("nvidia-smi") is not None


def _hardware_flags(labels: dict, no_gpu: bool, no_usb: bool, no_host_net: bool) -> tuple[list[str], list[str]]:
    flags: list[str] = []
    notes: list[str] = []
    profile = (labels.get(_LABEL_PROFILE) or "").lower()
    hardware = (labels.get(_LABEL_HARDWARE) or "").lower()
    hw_items = {h.strip() for h in hardware.split(",") if h.strip()}

    wants_gpu = any(k in profile for k in ("nvidia", "holoscan", "cuda"))
    if wants_gpu and not no_gpu:
        if _gpu_available():
            flags += ["--gpus", "all"]
        else:
            notes.append("image wants GPU but no NVIDIA runtime detected — skipping --gpus (use --force-gpu to override)")

    if hw_items & {"pluto", "rtlsdr", "hackrf", "bladerf"} and not no_usb:
        flags += ["--device", "/dev/bus/usb"]

    if hw_items & {"usrp", "thinkrf", "pluto"} and not no_host_net:
        flags += ["--net", "host"]

    return flags, notes


def _cmd_configure(args: argparse.Namespace) -> int:
    cfg = _config.load()
    if args.registry:
        cfg.registry = args.registry
    if args.namespace:
        cfg.namespace = args.namespace
    if args.sudo is not None:
        cfg.sudo = args.sudo
    path = _config.save(cfg)
    print(f"Saved app config to {path}")
    print(f"  registry:  {cfg.registry or '(unset)'}")
    print(f"  namespace: {cfg.namespace or '(unset)'}")
    print(f"  sudo:      {cfg.sudo}")
    return 0


def _cmd_pull(args: argparse.Namespace) -> int:
    cfg = _config.load()
    engine = _engine(cfg, args.sudo)
    ref = _resolve_ref(args.app, cfg)
    print(f"Pulling {ref}")
    return subprocess.call([*engine, "pull", ref])


def _cmd_run(args: argparse.Namespace) -> int:
    cfg = _config.load()
    engine = _engine(cfg, args.sudo)
    ref = _resolve_ref(args.app, cfg)

    if not _inspect_labels(engine, ref):
        rc = subprocess.call([*engine, "pull", ref])
        if rc != 0:
            return rc

    labels = _inspect_labels(engine, ref)
    no_gpu = args.no_gpu and not args.force_gpu
    hw_flags, notes = _hardware_flags(labels, no_gpu=no_gpu, no_usb=args.no_usb, no_host_net=args.no_host_net)
    if args.force_gpu and "--gpus" not in hw_flags:
        hw_flags = ["--gpus", "all", *hw_flags]

    cmd = [*engine, "run", "--rm"]
    if not args.foreground:
        cmd += ["-d"]
    cmd += ["--name", args.name or _container_name(ref)]
    cmd += hw_flags

    if args.config:
        cmd += ["-v", f"{args.config}:/config/config.yaml:ro", "-e", "RIA_CONFIG=/config/config.yaml"]

    for env in args.env or []:
        cmd += ["-e", env]
    for vol in args.volume or []:
        cmd += ["-v", vol]
    for port in args.publish or []:
        cmd += ["-p", port]

    cmd += list(args.docker_args or [])
    cmd += [ref]
    cmd += list(args.app_args or [])

    if args.dry_run:
        print(" ".join(cmd))
        return 0

    label_str = ", ".join(f"{k}={v}" for k, v in labels.items() if k.startswith("ria.")) or "(no ria.* labels)"
    print(f"Running {ref} [{label_str}]")
    if hw_flags:
        print(f"  auto flags: {' '.join(hw_flags)}")
    for note in notes:
        print(f"  note: {note}")
    return subprocess.call(cmd)


def _cmd_list(args: argparse.Namespace) -> int:
    cfg = _config.load()
    engine = _engine(cfg, args.sudo)
    return subprocess.call(
        [
            *engine,
            "images",
            "--filter",
            f"label={_LABEL_APP}",
            "--format",
            "table {{.Repository}}:{{.Tag}}\t{{.ID}}\t{{.Size}}",
        ]
    )


def _cmd_stop(args: argparse.Namespace) -> int:
    cfg = _config.load()
    engine = _engine(cfg, args.sudo)
    name = args.name or _container_name(_resolve_ref(args.app, cfg))
    return subprocess.call([*engine, "stop", name])


def _cmd_logs(args: argparse.Namespace) -> int:
    cfg = _config.load()
    engine = _engine(cfg, args.sudo)
    name = args.name or _container_name(_resolve_ref(args.app, cfg))
    cmd = [*engine, "logs"]
    if args.follow:
        cmd += ["-f"]
    cmd += [name]
    return subprocess.call(cmd)


def main() -> None:
    parser = argparse.ArgumentParser(prog="ria-app")
    parser.add_argument("--sudo", action="store_true", default=False, help="Run docker/podman via sudo")
    sub = parser.add_subparsers(dest="command", required=True)

    p_cfg = sub.add_parser("configure", help="Set default registry/namespace")
    p_cfg.add_argument("--registry", default=None, help="Default container registry (e.g. registry.riahub.ai)")
    p_cfg.add_argument("--namespace", default=None, help="Default namespace (e.g. qoherent)")
    p_cfg.add_argument(
        "--sudo",
        dest="sudo",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="Persist sudo default (--sudo / --no-sudo)",
    )

    p_pull = sub.add_parser("pull", help="Pull an app image")
    p_pull.add_argument("app", help="App name or image reference")

    p_run = sub.add_parser("run", help="Run an app, auto-detecting hardware flags")
    p_run.add_argument("app", help="App name or image reference")
    p_run.add_argument("--name", default=None, help="Container name (default: ria-app-<app>)")
    p_run.add_argument("--config", default=None, help="Path to config.yaml to mount into the container")
    p_run.add_argument("-e", "--env", action="append", help="Extra env var (KEY=VALUE)")
    p_run.add_argument("-v", "--volume", action="append", help="Extra volume mount")
    p_run.add_argument("-p", "--publish", action="append", help="Publish port")
    p_run.add_argument("--foreground", "-F", action="store_true", help="Run in foreground (no -d)")
    p_run.add_argument("--no-gpu", action="store_true", help="Skip --gpus flag even if image wants GPU")
    p_run.add_argument("--force-gpu", action="store_true", help="Force --gpus all even if no NVIDIA runtime detected")
    p_run.add_argument("--no-usb", action="store_true", help="Skip --device /dev/bus/usb")
    p_run.add_argument("--no-host-net", action="store_true", help="Skip --net host")
    p_run.add_argument("--dry-run", action="store_true", help="Print the container command and exit")
    p_run.add_argument("--docker-args", nargs=argparse.REMAINDER, help="Pass remaining args to docker/podman run")
    p_run.add_argument("--app-args", nargs=argparse.REMAINDER, help="Pass remaining args to the app entrypoint")

    sub.add_parser("list", help="List locally cached RIA app images")

    p_stop = sub.add_parser("stop", help="Stop a running app")
    p_stop.add_argument("app", help="App name or image reference")
    p_stop.add_argument("--name", default=None, help="Container name override")

    p_logs = sub.add_parser("logs", help="Tail logs of a running app")
    p_logs.add_argument("app", help="App name or image reference")
    p_logs.add_argument("--name", default=None, help="Container name override")
    p_logs.add_argument("-f", "--follow", action="store_true", help="Follow log output")

    args = parser.parse_args()

    dispatch = {
        "configure": _cmd_configure,
        "pull": _cmd_pull,
        "run": _cmd_run,
        "list": _cmd_list,
        "stop": _cmd_stop,
        "logs": _cmd_logs,
    }
    sys.exit(dispatch[args.command](args))


if __name__ == "__main__":
    main()