manifold-cli 0.0.39__py3-none-any.whl

manifold_cli/commands/benchmark.py

@@ -0,0 +1,539 @@
+"""The `manifold benchmark` command group."""
+
+import asyncio
+import json
+import os
+from typing import Annotated
+
+import typer
+from rich.console import Console
+from rich.markup import escape
+from rich.table import Table
+from rich.text import Text
+
+from manifold_cli.bridge import run_bridge
+from manifold_cli.clients.image import ImageInspectError
+from manifold_cli.commands.organization import (
+    ORGANIZATION_CALL_ERRORS,
+    exit_for_call_error,
+    resolve_owner,
+)
+from manifold_cli.commands.table import catalog_table, print_table, terminal_width
+from manifold_cli.dependencies import Dependencies
+from manifold_cli.models.benchmark import (
+    BenchmarkNotFoundError,
+    BenchmarkVersion,
+    InvalidImageRubricError,
+    InvalidImageSpecError,
+)
+from manifold_cli.models.catalog import (
+    CatalogRegisterError,
+    InvalidImageError,
+    InvalidReferenceError,
+    Visibility,
+    join_versioned_ref,
+    split_ref,
+    split_versioned_ref,
+)
+from manifold_cli.models.run import SchedulingType
+from manifold_cli.transport import DEFAULT_TRANSPORTS, TcpTransport, TransportError
+from manifold_cli.transport.diagnostics import PingResult, ping
+
+DEFAULT_BRIDGE_LISTEN = "127.0.0.1:9001"
+
+DEFAULT_PING_BYTES = 4 * 1024
+DEFAULT_PING_TIMEOUT_SECONDS = 30.0
+
+
+def _publish_rubric(
+    console: Console,
+    deps: Dependencies,
+    slug: str,
+    version: str,
+    owner: str,
+    *,
+    image: str,
+) -> None:
+    """Publish the rubric the just-registered image carries, and say so.
+
+    Silent for an image carrying none. An image this box does not have is
+    reported as unread rather than as carrying nothing.
+    """
+    try:
+        tasks = deps.benchmark_service.publish_image_rubric(
+            slug, version, owner, image=image
+        )
+    except ImageInspectError as exc:
+        console.print(f"[dim]Rubric not read: {escape(str(exc))}.[/dim]")
+        return
+    except InvalidImageRubricError as exc:
+        console.print(f"[red]{escape(str(exc))}[/red]")
+        raise typer.Exit(1) from exc
+    except (CatalogRegisterError, BenchmarkNotFoundError) as exc:
+        console.print(f"[red]{escape(str(exc))}[/red]")
+        raise typer.Exit(1) from exc
+    except ORGANIZATION_CALL_ERRORS as exc:
+        raise exit_for_call_error(console, exc) from exc
+    if tasks is None:
+        return
+    console.print(
+        f"[green]✓ Published rubric[/green] "
+        f"[dim]{tasks} task{'' if tasks == 1 else 's'}[/dim]"
+    )
+
+
+def _read_spec(console: Console, deps: Dependencies, image: str) -> dict | None:
+    """The spec read from `image`, or None when the image has no spec file or
+    cannot be read.
+
+    An unreadable image registers without a spec, and the output says so. A
+    spec that is not a JSON object stops the command before it registers
+    anything.
+    """
+    try:
+        return deps.benchmark_service.read_image_spec(image)
+    except ImageInspectError as exc:
+        console.print(f"[dim]Spec not read: {escape(str(exc))}.[/dim]")
+        return None
+    except InvalidImageSpecError as exc:
+        console.print(f"[red]{escape(str(exc))}[/red]")
+        raise typer.Exit(1) from exc
+
+
+def _describe_shards(version: BenchmarkVersion) -> str:
+    if version.shards_per_gpu is not None:
+        return f"{version.shards_per_gpu} shards per GPU"
+    return f"{version.shards} shard{'' if version.shards == 1 else 's'}"
+
+
+def register(app: typer.Typer, deps: Dependencies) -> None:
+    console = Console()
+    benchmark_app = typer.Typer(
+        help="Run benchmarks against a policy.", no_args_is_help=True
+    )
+
+    @benchmark_app.command("list")
+    def list_benchmarks() -> None:
+        """List the benchmarks you can see."""
+        try:
+            benchmarks = deps.benchmark_service.list_benchmarks()
+        except ORGANIZATION_CALL_ERRORS as exc:
+            raise exit_for_call_error(console, exc) from exc
+        if len(benchmarks) == 0:
+            console.print("There are no benchmarks yet.")
+            return
+        print_table(console, catalog_table(benchmarks, terminal_width(console)))
+
+    @benchmark_app.command()
+    def versions(
+        benchmark: Annotated[
+            str,
+            typer.Argument(
+                help="Benchmark identifier, or <org>/<identifier> to specify its owner."
+            ),
+        ],
+    ) -> None:
+        """List the versions of a benchmark, newest first."""
+        try:
+            owner, slug = split_ref(benchmark)
+        except InvalidReferenceError as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        # A qualified reference names its owner outright. A bare slug resolves
+        # the way `run submit` resolves one: the caller's own benchmark first,
+        # then a public benchmark of that slug.
+        list_versions = (
+            deps.benchmark_service.list_versions
+            if owner is not None
+            else deps.benchmark_service.list_own_or_public_versions
+        )
+        if owner is None:
+            try:
+                owner = resolve_owner(console, deps.organization_service)
+            except ORGANIZATION_CALL_ERRORS as exc:
+                raise exit_for_call_error(console, exc) from exc
+        try:
+            found = list_versions(slug, owner)
+        except BenchmarkNotFoundError as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        except ORGANIZATION_CALL_ERRORS as exc:
+            raise exit_for_call_error(console, exc) from exc
+        if len(found) == 0:
+            console.print(f"'{escape(benchmark)}' has no versions yet.")
+            return
+        table = Table(box=None, pad_edge=False)
+        table.add_column("Version")
+        table.add_column("Adapter", style="dim")
+        table.add_column("Episodes", style="dim", justify="right")
+        table.add_column("Shards", style="dim", justify="right")
+        table.add_column("Shards/GPU", style="dim", justify="right")
+        for version in found:
+            table.add_row(
+                Text(version.version, style="bold"),
+                version.adapter.value,
+                Text(str(version.episodes)),
+                Text(str(version.shards)),
+                Text(
+                    str(version.shards_per_gpu)
+                    if version.shards_per_gpu is not None
+                    else "-"
+                ),
+            )
+        print_table(console, table)
+
+    @benchmark_app.command()
+    def init(
+        benchmark: Annotated[
+            str,
+            typer.Argument(
+                help="Benchmark identifier to register the image under, "
+                "or <org>/<identifier> to specify the organization that owns it."
+            ),
+        ],
+        image: Annotated[
+            str,
+            typer.Option(
+                "--image",
+                help="Container image, e.g. ghcr.io/bifrost/libero-spatial:0.1.0. "
+                "Its tag is the version this registers.",
+            ),
+        ],
+        episodes: Annotated[
+            int,
+            typer.Option(
+                "--episodes",
+                min=1,
+                help="Episodes this image runs. Every run of the version "
+                "executes that many.",
+            ),
+        ],
+        minimum_gpu_memory_gb: Annotated[
+            int,
+            typer.Option(
+                "--minimum-gpu-memory-gb",
+                min=0,
+                help="Minimum memory required on one GPU. Set to 0 to allow "
+                "placement on runners without a GPU.",
+            ),
+        ],
+        minimum_driver_major: Annotated[
+            int | None,
+            typer.Option(
+                "--minimum-driver-major",
+                min=1,
+                help="Minimum NVIDIA driver major version required on a "
+                "runner. Omit to leave the lower bound open.",
+            ),
+        ] = None,
+        maximum_driver_major: Annotated[
+            int | None,
+            typer.Option(
+                "--maximum-driver-major",
+                min=1,
+                help="Maximum NVIDIA driver major version allowed on a "
+                "runner. Omit to leave the upper bound open.",
+            ),
+        ] = None,
+        minimum_gpu_count: Annotated[
+            int | None,
+            typer.Option(
+                "--minimum-gpu-count",
+                min=0,
+                help="Minimum number of GPUs required on one runner. If omitted, "
+                "the default is 1, or 0 when --minimum-gpu-memory-gb is 0.",
+            ),
+        ] = None,
+        shards: Annotated[
+            int,
+            typer.Option(
+                "--shards",
+                min=1,
+                help="Parallel shards the image partitions its episodes across. "
+                "A run fans out to this many unless `run submit --shards` names "
+                "another count; 1 does not partition. A count above `--episodes` "
+                "is clamped to it at registration.",
+            ),
+        ] = 1,
+        shards_per_gpu: Annotated[
+            int | None,
+            typer.Option(
+                "--shards-per-gpu",
+                min=1,
+                help="Parallel shards per GPU of each benchmark runner, in place "
+                "of --shards. A run fans out to this many per GPU, up to one shard "
+                "per episode, unless `run submit` sets another count.",
+            ),
+        ] = None,
+        scheduling_type: Annotated[
+            SchedulingType,
+            typer.Option(
+                "--scheduling-type",
+                help="How the benchmark side divides a run's episodes among its "
+                "shards: with `preassigned`, each shard's worker runs the episodes "
+                "of its own shard, and with `task_queue`, each worker runs tasks "
+                "claimed from the run's task queue.",
+            ),
+        ] = SchedulingType.PREASSIGNED,
+        allow_preemptible: Annotated[
+            bool,
+            typer.Option(
+                "--allow-preemptible",
+                help="Let a run of this version use runners whose VM the "
+                "provider may stop at any time, for both its policy and its "
+                "benchmark side.",
+            ),
+        ] = False,
+        display_name: Annotated[
+            str | None,
+            typer.Option(
+                "--display-name",
+                help="Name shown in `benchmark list`. Left off, an existing "
+                "benchmark keeps the name it has and a new one uses its identifier.",
+            ),
+        ] = None,
+        visibility: Annotated[
+            Visibility | None,
+            typer.Option(
+                "--visibility",
+                help="Who can read the benchmark. Left off, an existing "
+                "benchmark keeps the visibility it has and a new one is visible "
+                "to its organization.",
+            ),
+        ] = None,
+    ) -> None:
+        """Register a container image as a version of a benchmark."""
+        try:
+            owner, slug = split_ref(benchmark)
+        except InvalidReferenceError as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        spec = _read_spec(console, deps, image)
+        try:
+            if owner is None:
+                owner = resolve_owner(console, deps.organization_service)
+            version = deps.benchmark_service.register_container_version(
+                slug,
+                owner,
+                image=image,
+                episodes=episodes,
+                minimum_gpu_memory_gb=minimum_gpu_memory_gb,
+                minimum_driver_major=minimum_driver_major,
+                maximum_driver_major=maximum_driver_major,
+                minimum_gpu_count=minimum_gpu_count,
+                scheduling_type=scheduling_type,
+                allow_preemptible=allow_preemptible,
+                shards=shards,
+                shards_per_gpu=shards_per_gpu,
+                display_name=display_name,
+                visibility=visibility,
+                spec=spec,
+            )
+        except (
+            InvalidImageError,
+            CatalogRegisterError,
+            BenchmarkNotFoundError,
+        ) as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        except ORGANIZATION_CALL_ERRORS as exc:
+            raise exit_for_call_error(console, exc) from exc
+        console.print(
+            f"[green]✓ Registered[/green] [bold]{escape(benchmark)}[/bold] "
+            f"[bold]{escape(version.version)}[/bold] "
+            f"[dim]{version.adapter.value}, {version.episodes} episodes, "
+            f"{_describe_shards(version)}[/dim]"
+        )
+        console.print(f"  {escape(image)}")
+        if spec is not None:
+            console.print("[green]✓ Recorded spec[/green]")
+        _publish_rubric(console, deps, slug, version.version, owner, image=image)
+        console.print(
+            f"Run it with: manifold run submit <policy> {escape(benchmark)} --name NAME"
+        )
+
+    @benchmark_app.command()
+    def spec(
+        benchmark: Annotated[
+            str,
+            typer.Argument(
+                help="Benchmark version, as [<org>/]<identifier>:<version>."
+            ),
+        ],
+    ) -> None:
+        """Print a benchmark version's spec as JSON.
+
+        The spec lists the observations, actions and embodiment the benchmark
+        sends to a policy. It was read from the image when the version was
+        registered. Save the output to pair a policy with the benchmark.
+        """
+        try:
+            owner, slug, version = split_versioned_ref(benchmark)
+        except InvalidReferenceError as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        if version is None:
+            console.print(
+                f"[red]'{escape(benchmark)}' names no version; "
+                "write <identifier>:<version>.[/red]"
+            )
+            raise typer.Exit(1)
+        or_public = owner is None
+        try:
+            if owner is None:
+                owner = resolve_owner(console, deps.organization_service)
+            found = deps.benchmark_service.version_spec(
+                slug, owner, version, or_public=or_public
+            )
+        except BenchmarkNotFoundError as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        except ORGANIZATION_CALL_ERRORS as exc:
+            raise exit_for_call_error(console, exc) from exc
+        if found is None:
+            console.print(
+                f"[red]'{escape(benchmark)}' was registered without a spec.[/red]"
+            )
+            raise typer.Exit(1)
+        typer.echo(json.dumps(found, indent=2))
+
+    @benchmark_app.command()
+    def delete(
+        benchmark: Annotated[
+            str,
+            typer.Argument(
+                help="Benchmark to delete, as [<org>/]<identifier>[:<version>], for "
+                "example libero-spatial:0.0.1. With a version, only that "
+                "version is deleted."
+            ),
+        ],
+    ) -> None:
+        """Delete a benchmark, or one version of it.
+
+        A deleted benchmark or version is left out of listings and cannot be used
+        in new runs. `run get` still shows it on the runs that used it.
+        Registering it again restores it. A restored benchmark takes the display
+        name and visibility of that registration, not its earlier ones.
+        """
+        try:
+            owner, slug, version = split_versioned_ref(benchmark)
+        except InvalidReferenceError as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        try:
+            if owner is None:
+                owner = resolve_owner(console, deps.organization_service)
+            deps.benchmark_service.delete(slug, owner, version)
+        except BenchmarkNotFoundError as exc:
+            console.print(f"[red]{escape(str(exc))}[/red]")
+            raise typer.Exit(1) from exc
+        except ORGANIZATION_CALL_ERRORS as exc:
+            raise exit_for_call_error(console, exc) from exc
+        deleted = (
+            "Benchmark deleted" if version is None else "Benchmark version deleted"
+        )
+        console.print(
+            f"[green]✓ {deleted}[/green] "
+            f"[bold]{escape(join_versioned_ref(owner, slug, version))}[/bold]"
+        )
+
+    @benchmark_app.command(hidden=True)
+    def run(
+        target: Annotated[
+            str,
+            typer.Argument(help="Policy target, e.g. tcp:127.0.0.1:9000 or iroh:ID."),
+        ],
+        use_ping: Annotated[
+            bool,
+            typer.Option(
+                "--ping", help="Run a built-in ping benchmark (round-trips a payload)."
+            ),
+        ] = False,
+        size: Annotated[
+            int,
+            typer.Option("--size", help="Ping payload size in bytes."),
+        ] = DEFAULT_PING_BYTES,
+        timeout: Annotated[
+            float,
+            typer.Option(
+                "--timeout",
+                help="Abort the ping if it has not completed after this many seconds.",
+            ),
+        ] = DEFAULT_PING_TIMEOUT_SECONDS,
+    ) -> None:
+        """Run a benchmark against the policy at the given transport target."""
+        if not use_ping:
+            console.print(
+                "[red]Real benchmarks are not available yet.[/red] "
+                "Pass --ping to run the built-in ping benchmark."
+            )
+            raise typer.Exit(1)
+        payload = os.urandom(size)
+
+        async def run_ping() -> PingResult:
+            async with asyncio.timeout(timeout):
+                return await ping(connector, payload)
+
+        try:
+            transport, rest = DEFAULT_TRANSPORTS.resolve(target)
+            connector = transport.build_connector(rest)
+            result = asyncio.run(run_ping())
+        except TimeoutError as exc:
+            console.print(f"[red]Ping timed out[/red] after {timeout:g} s.")
+            raise typer.Exit(1) from exc
+        except (TransportError, OSError) as exc:
+            console.print(f"[red]Ping failed:[/red] {escape(str(exc))}")
+            raise typer.Exit(1) from exc
+
+        if not result.matched:
+            console.print(
+                f"[red]Mismatch:[/red] sent {result.size} bytes, echo differed."
+            )
+            raise typer.Exit(1)
+
+        console.print(
+            f"[green]OK[/green] round-tripped {result.size} bytes via "
+            f"[bold]{escape(target)}[/bold] in [bold]{result.elapsed_ms:.2f} ms[/bold]."
+        )
+
+    @benchmark_app.command(hidden=True)
+    def bridge(
+        target: Annotated[
+            str,
+            typer.Argument(
+                help="Policy-bridge target for our transport, e.g. "
+                "tcp:100.101.132.68:9500 or iroh:ENDPOINT_ID."
+            ),
+        ],
+        listen: Annotated[
+            str,
+            typer.Option(
+                "--listen", help="Local address to accept the SDK benchmark, HOST:PORT."
+            ),
+        ] = DEFAULT_BRIDGE_LISTEN,
+    ) -> None:
+        """Expose a local TCP socket forwarding to a policy bridge over transport."""
+
+        async def run() -> None:
+            transport, rest = DEFAULT_TRANSPORTS.resolve(target)
+            connector = transport.build_connector(rest)
+            listener = await TcpTransport().open_listener(listen)
+            console.print(
+                "[green]Benchmark bridge up.[/green] Point the SDK benchmark's "
+                f"--server at: [bold]{escape(listener.target)}[/bold]"
+            )
+            console.print(
+                f"Forwarding to policy bridge at [bold]{escape(target)}[/bold]."
+            )
+            console.print("Bridging... (press Ctrl-C to stop)")
+            await run_bridge(listener, connector)
+
+        try:
+            asyncio.run(run())
+        except KeyboardInterrupt:
+            console.print("\n[yellow]Benchmark bridge stopped.[/yellow]")
+        except (TransportError, OSError) as exc:
+            console.print(f"[red]Bridge failed:[/red] {escape(str(exc))}")
+            raise typer.Exit(1) from exc
+
+    app.add_typer(benchmark_app, name="benchmark")

manifold_cli/commands/organization.py

@@ -0,0 +1,154 @@
+from http import HTTPStatus
+from json import JSONDecodeError
+
+import httpx
+import typer
+from pydantic import ValidationError
+from rich.console import Console
+from rich.markup import escape
+from rich.table import Table
+from rich.text import Text
+
+from manifold_cli.commands.table import print_table
+from manifold_cli.dependencies import Dependencies
+from manifold_cli.models.errors import (
+    ManifoldApiError,
+    NotLoggedInError,
+    SessionExpiredError,
+)
+from manifold_cli.services.organization import OrganizationService
+
+NO_ORGANIZATION_MESSAGE = (
+    "You do not have an organization yet. Reach out to Bifrost to get set up."
+)
+
+# Failures an organization API call can raise: no/expired credentials, a
+# transport or HTTP-status error, or a 2xx body that isn't the JSON shape we
+# expect. The login epilogue in commands.auth reuses this to stay in lockstep.
+ORGANIZATION_CALL_ERRORS = (
+    NotLoggedInError,
+    SessionExpiredError,
+    ManifoldApiError,
+    httpx.HTTPError,
+    JSONDecodeError,  # a 2xx body that isn't JSON
+    ValidationError,  # a 2xx body that doesn't match our schema
+)
+
+
+def resolve_organization(
+    console: Console,
+    service: OrganizationService,
+    organization: str | None,
+) -> str:
+    """Settle which organization a --org command acts as.
+
+    Every policy, benchmark and run call is made as one organization. Anyone
+    who belongs to multiple organizations must specify which to use.
+    """
+    if organization is not None:
+        return organization.strip().lower()
+    organizations = service.list()
+    if len(organizations) == 1:
+        return organizations[0].slug
+    if len(organizations) == 0:
+        console.print(NO_ORGANIZATION_MESSAGE)
+        raise typer.Exit(1)
+    slugs = ", ".join(o.slug for o in organizations)
+    console.print(
+        f"You belong to several organizations ({slugs}); pick one with --org."
+    )
+    raise typer.Exit(1)
+
+
+def resolve_owner(console: Console, service: OrganizationService) -> str:
+    """Settle who owns the entity a bare `<slug>` reference names.
+
+    A qualified `<org>/<slug>` names its owner already, so this is the bare
+    slug's fallback: the caller's own organization, or the qualified form when
+    they belong to several.
+    """
+    organizations = service.list()
+    if len(organizations) == 1:
+        return organizations[0].slug
+    if len(organizations) == 0:
+        console.print(NO_ORGANIZATION_MESSAGE)
+        raise typer.Exit(1)
+    slugs = ", ".join(o.slug for o in organizations)
+    console.print(
+        f"You belong to several organizations ({slugs}); "
+        "specify the owner as <org>/<identifier>."
+    )
+    raise typer.Exit(1)
+
+
+def exit_for_call_error(console: Console, exc: Exception) -> typer.Exit:
+    """Report a failed API call, and return the Exit the command raises.
+
+    Every command group shares this so each failure names its own cause: a
+    missing or lapsed session names the command that fixes it, an HTTP error
+    status is reported with what the server said, and only a transport failure
+    is reported as Manifold being unreachable.
+    """
+    match exc:
+        case NotLoggedInError():
+            console.print("Not logged in. Run `manifold auth login` to authenticate.")
+        case SessionExpiredError():
+            console.print(
+                "[yellow]Session expired.[/yellow] "
+                "Run `manifold auth login` to re-authenticate."
+            )
+        case ManifoldApiError():
+            console.print(_api_error_message(exc))
+        case JSONDecodeError() | ValidationError():
+            console.print(
+                f"[red]Manifold returned an unexpected response:[/red] "
+                f"{escape(str(exc))}"
+            )
+        case _:
+            console.print(f"[red]Could not reach Manifold:[/red] {escape(str(exc))}")
+    return typer.Exit(1)
+
+
+def _api_error_message(exc: ManifoldApiError) -> str:
+    """The line for an HTTP error status no domain exception claimed.
+
+    The server's `detail` is the message where it sent one, and an excerpt of
+    the raw body is the fallback, clipped to the 200 characters `str(exc)`
+    shows.
+    """
+    said = escape(exc.detail or exc.body[:200]) or "no detail given"
+    if exc.status_code == HTTPStatus.FORBIDDEN:
+        return f"[red]You do not have permission for this:[/red] {said}"
+    if exc.status_code >= HTTPStatus.INTERNAL_SERVER_ERROR:
+        return f"[red]Manifold failed (HTTP {exc.status_code}):[/red] {said}"
+    return f"[red]Manifold rejected the request (HTTP {exc.status_code}):[/red] {said}"
+
+
+def register(app: typer.Typer, deps: Dependencies) -> None:
+    console = Console()
+    org_app = typer.Typer(
+        name="org",
+        help="Manage organizations.",
+        no_args_is_help=True,
+    )
+    app.add_typer(org_app)
+
+    @org_app.command("list")
+    def list_organizations() -> None:
+        """List the organizations you belong to."""
+        try:
+            organizations = deps.organization_service.list()
+        except ORGANIZATION_CALL_ERRORS as exc:
+            raise exit_for_call_error(console, exc) from exc
+        if not organizations:
+            console.print(NO_ORGANIZATION_MESSAGE)
+            return
+        table = Table(box=None, pad_edge=False)
+        table.add_column("Identifier")
+        table.add_column("Name")
+        for organization in organizations:
+            table.add_row(
+                Text(organization.slug, style="bold"),
+                Text(organization.display_name),
+            )
+        print_table(console, table)