public_sdk_launch.py

#!/usr/bin/env python3
"""Launch and monitor a MirrorNeuron job with the limited public Python SDK.

Examples:
    python examples/public_sdk_launch.py /absolute/path/to/bundle \
        --inputs '{"query":"example"}'

    python examples/public_sdk_launch.py --job-id job_example-12345678 \
        --inputs '{"query":"another"}'

The target and connection token are discovered from local MirrorNeuron config
unless --target or --token-file is supplied.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

from mn_sdk.public import LaunchError, MirrorNeuron, RunTimeoutError


def json_object(text: str, *, source: str) -> dict[str, Any]:
    try:
        value = json.loads(text)
    except json.JSONDecodeError as exc:
        raise ValueError(f"{source} must contain valid JSON: {exc}") from exc
    if not isinstance(value, dict):
        raise TypeError(f"{source} must contain a JSON object")
    return value


def load_inputs(args: argparse.Namespace) -> dict[str, Any]:
    if args.inputs_file is not None:
        return json_object(
            args.inputs_file.read_text(encoding="utf-8"),
            source=str(args.inputs_file),
        )
    return json_object(args.inputs, source="--inputs")


def load_connection_token(path: Path | None) -> str | None:
    if path is None:
        return None
    token = path.read_text(encoding="utf-8").strip()
    if not token:
        raise ValueError(f"connection token file is empty: {path}")
    return token


def run_status(record: dict[str, Any]) -> str:
    nested = record.get("run") if isinstance(record.get("run"), dict) else {}
    return str(record.get("status") or nested.get("status") or "unknown")


def print_json(label: str, value: dict[str, Any]) -> None:
    print(f"{label}: {json.dumps(value, sort_keys=True, default=str)}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Launch a MirrorNeuron bundle or start another run for an existing "
            "stable job, then monitor it."
        )
    )
    parser.add_argument(
        "bundle",
        type=Path,
        nargs="?",
        help="Bundle directory containing manifest.json.",
    )
    parser.add_argument(
        "--job-id",
        help="Start a new run for this existing stable job instead of launching a bundle.",
    )
    input_group = parser.add_mutually_exclusive_group()
    input_group.add_argument(
        "--inputs",
        default="{}",
        help="Run inputs as a JSON object (default: {}).",
    )
    input_group.add_argument(
        "--inputs-file",
        type=Path,
        help="Path to a UTF-8 JSON object containing run inputs.",
    )
    parser.add_argument(
        "--target",
        help="Optional Core gRPC host:port; local config is used when omitted.",
    )
    parser.add_argument(
        "--token-file",
        type=Path,
        help="Optional connection-token file; local config is used when omitted.",
    )
    parser.add_argument(
        "--run-id",
        help="Optional caller-supplied run identity.",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=600.0,
        help="Seconds for wait() after event streaming finishes (default: 600).",
    )
    parser.add_argument(
        "--poll-interval",
        type=float,
        default=1.0,
        help="Status polling interval in seconds (default: 1).",
    )
    parser.add_argument(
        "--no-events",
        action="store_true",
        help="Skip the following event stream and monitor only with wait().",
    )
    parser.add_argument(
        "--include-heartbeats",
        action="store_true",
        help="Print stream heartbeat events.",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Force bundle launch by skipping supported soft validation checks.",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if (args.bundle is None) == (args.job_id is None):
        parser.error("pass exactly one bundle path or --job-id")
    if args.force and args.job_id is not None:
        parser.error("--force applies only when launching a bundle")

    try:
        inputs = load_inputs(args)
        connection_token = load_connection_token(args.token_file)
        with MirrorNeuron(
            target=args.target,
            connection_token=connection_token,
        ) as mn:
            if args.job_id is not None:
                run = mn.start(
                    args.job_id,
                    inputs=inputs,
                    run_id=args.run_id,
                )
            else:
                run = mn.launch(
                    args.bundle,
                    inputs=inputs,
                    run_id=args.run_id,
                    force=args.force,
                )

            print(f"target: {mn.target}")
            print(f"job_id: {run.job_id}")
            print(f"run_id: {run.run_id}")
            print_json("accepted", run.initial)

            if not args.no_events:
                for event in run.events(
                    include_heartbeats=args.include_heartbeats,
                ):
                    print_json("event", event)

            final = run.wait(
                timeout=args.timeout,
                poll_interval=args.poll_interval,
            )
            print_json("final", final)
            return 0 if run_status(final).lower() == "completed" else 1
    except LaunchError as exc:
        print(
            "launch failed: "
            f"stage={exc.stage} job_id={exc.job_id} run_id={exc.run_id}: {exc}",
            file=sys.stderr,
        )
        if exc.report:
            print_json("validation", exc.report)
        return 1
    except RunTimeoutError as exc:
        print(
            f"run still active after {exc.timeout:g}s: "
            f"job_id={exc.job_id} run_id={exc.run_id}",
            file=sys.stderr,
        )
        return 2
    except (OSError, TypeError, ValueError) as exc:
        print(f"invalid input: {exc}", file=sys.stderr)
        return 2
    except KeyboardInterrupt:
        print("monitoring interrupted; the accepted run was not cancelled", file=sys.stderr)
        return 130


if __name__ == "__main__":
    raise SystemExit(main())