flowmesh-cli 0.1.0__py3-none-any.whl

flowmesh_cli/commands/ssh.py

@@ -0,0 +1,422 @@
+"""SSH session commands."""
+
+import asyncio
+import json
+import os
+import shlex
+import shutil
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import typer
+import websockets
+from flowmesh import FlowMesh
+from flowmesh.exceptions import FlowMeshError, NotFoundError
+from flowmesh.models.common import TERMINAL_TASK_STATUSES, TaskStatus
+from flowmesh.params import append_param, extend_params
+
+from ..core import logging
+from ..core.query import parse_query_filters
+from ..core.typer import get_typer
+
+app = get_typer(help="Connect to SSH sessions on FlowMesh workers.")
+
+
+def _exec_ssh(
+    ssh_info: dict[str, Any],
+    task_id: str,
+    extra_args: str | None,
+    direct: bool = False,
+) -> None:
+    """Replace the current process with an ssh command."""
+    ssh_bin = shutil.which("ssh")
+    if ssh_bin is None:
+        logging.error("ssh not found in PATH")
+        raise typer.Exit(code=1)
+
+    mode = str(ssh_info.get("mode", "direct"))
+    user = ssh_info.get("username", "flowmesh")
+    host = ssh_info.get("host")
+    port = ssh_info.get("port")
+    if direct:
+        if direct_host := ssh_info.get("directHost"):
+            host = direct_host
+        if direct_port := ssh_info.get("directPort"):
+            port = direct_port
+
+    args = [
+        ssh_bin,
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "UserKnownHostsFile=/dev/null",
+        "-o",
+        "LogLevel=ERROR",
+    ]
+
+    if direct or mode in ("direct", "forward"):
+        if not host:
+            logging.error("SSH host is not available.")
+            raise typer.Exit(code=1)
+        if port:
+            args += ["-p", str(port)]
+        args.append(f"{user}@{host}")
+    elif mode == "proxy":
+        proxy_cmd = f"flowmesh ssh proxy {task_id}"
+        args += ["-o", f"ProxyCommand={proxy_cmd}"]
+        args.append(f"{user}@{task_id}")
+    else:
+        logging.error(f"Unknown SSH publish mode: {mode}")
+        raise typer.Exit(code=1)
+
+    if extra_args:
+        args.extend(extra_args.split())
+
+    logging.info(f"Connecting: {' '.join(args)}")
+    os.execvp(ssh_bin, args)
+
+
+def _exit_for_task_status(task_id: str, task_info: Any) -> None:
+    status = task_info.status
+    if status == TaskStatus.DONE:
+        logging.info(f"Task {task_id} completed successfully.")
+        raise typer.Exit(code=0)
+    if status == TaskStatus.FAILED:
+        logging.error(f"Task {task_id} failed: {task_info.error or 'unknown error'}")
+        raise typer.Exit(code=1)
+    if status == TaskStatus.CANCELLED:
+        logging.info(f"Task {task_id} was cancelled.")
+        raise typer.Exit(code=130)
+
+
+def _poll_status_and_exit(client: FlowMesh, task_id: str, interval: float) -> None:
+    try:
+        task_info = client.tasks.wait(task_id, interval=interval)
+    except KeyboardInterrupt:
+        logging.warning("Interrupted.")
+        raise typer.Exit(code=130)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    _exit_for_task_status(task_id, task_info)
+
+
+def _stream_logs_and_exit(
+    client: FlowMesh,
+    task_id: str,
+    interval: float,
+    tail: bool = False,
+) -> None:
+    logging.info("Streaming task logs...")
+    cursor: str | None = "$" if tail else "0"
+    try:
+        while True:
+            saw_logs = False
+            try:
+                for entry in client.tasks.stream_logs(task_id, cursor=cursor):
+                    saw_logs = True
+                    cursor = entry.cursor or cursor
+                    event = entry.event.model_dump(mode="json")
+                    ts = event.get("ts", "")
+                    message = str(event.get("message", "")).rstrip("\n")
+                    prefix = f"[{ts}] " if ts else ""
+                    logging.log(prefix + message)
+            except NotFoundError:
+                pass
+
+            task_info = client.tasks.retrieve(task_id)
+            if task_info.status in TERMINAL_TASK_STATUSES:
+                _exit_for_task_status(task_id, task_info)
+            if not saw_logs:
+                time.sleep(interval)
+    except KeyboardInterrupt:
+        logging.warning("Interrupted.")
+        raise typer.Exit(code=130)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+
+
+# ------------------------------------------------------------------ #
+# Commands
+# ------------------------------------------------------------------ #
+
+
+@app.command("connect")
+def connect(
+    task_id: str = typer.Argument(..., help="Task identifier"),
+    interval: float = typer.Option(2.0, help="Polling interval in seconds"),
+    direct: bool = typer.Option(
+        False, "--direct", help="Use direct SSH instead of proxy/forward"
+    ),
+    ssh_args: str | None = typer.Option(
+        None, "--ssh-args", help="Extra arguments passed to ssh"
+    ),
+) -> None:
+    """Wait for an SSH session to be ready on a task and connect."""
+    logging.info(f"Waiting for SSH session on task {task_id}...")
+    client = FlowMesh()
+    try:
+        ssh_info = client.tasks.wait_for_ssh(task_id, interval=interval)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    _exec_ssh(ssh_info, task_id, ssh_args, direct)
+
+
+@app.command("run")
+def run(
+    key: Path | None = typer.Option(
+        None, "--key", "-k", help="Path to SSH public key file"
+    ),
+    gpu: int | None = typer.Option(None, "--gpu", "-g", help="Number of GPUs"),
+    gpu_memory: str | None = typer.Option(
+        None, "--gpu-memory", help="GPU memory requirement (e.g. 16GB)"
+    ),
+    memory: str | None = typer.Option(
+        None, "--memory", "-m", help="System memory (e.g. 8Gi)"
+    ),
+    cpu: int | None = typer.Option(None, "--cpu", help="Number of CPU cores"),
+    ttl: int = typer.Option(3600, "--ttl", help="Session TTL in seconds"),
+    idle_timeout: int = typer.Option(
+        900, "--idle-timeout", help="Idle timeout in seconds"
+    ),
+    image: str | None = typer.Option(None, "--image", help="Custom session image"),
+    user: str = typer.Option("flowmesh", "--user", "-u", help="SSH username"),
+    mode: str = typer.Option(
+        "proxy", "--mode", help="Publish mode: direct|proxy|forward"
+    ),
+    worker: str | None = typer.Option(
+        None, "--worker", "-w", help="Pin to a specific worker"
+    ),
+    name: str = typer.Option("ssh-session", "--name", "-n", help="Task name"),
+    env: list[str] | None = typer.Option(
+        None, "--env", "-e", help="Environment variable KEY=VALUE (repeatable)"
+    ),
+    command: str | None = typer.Option(
+        None,
+        "--command",
+        "-c",
+        help="Command to run non-interactively (e.g. 'python train.py')",
+    ),
+    entrypoint_override: str | None = typer.Option(
+        None,
+        "--entrypoint",
+        help="Override image entrypoint (non-interactive mode)",
+    ),
+    interactive: bool | None = typer.Option(
+        None,
+        "--interactive/--non-interactive",
+        help="Whether to use interactive SSH or non-interactive command mode.",
+    ),
+    logs: bool = typer.Option(
+        False,
+        "--logs/--no-logs",
+        help="For non-interactive runs, stream logs instead of only polling status.",
+    ),
+    tail: bool = typer.Option(
+        False,
+        "--tail",
+        help="With --logs, start from the latest log entry instead of the beginning.",
+    ),
+    interval: float = typer.Option(2.0, "--interval", help="Polling interval"),
+    direct: bool = typer.Option(
+        False, "--direct", help="Use direct SSH instead of proxy/forward"
+    ),
+    ssh_args: str | None = typer.Option(
+        None, "--ssh-args", help="Extra arguments passed to ssh"
+    ),
+) -> None:
+    """Submit an SSH task, wait for the session, and connect.
+
+    When --command or --entrypoint is provided, the task runs non-interactively:
+    the specified command executes in the container and the CLI polls until the
+    task completes.
+    """
+    inferred_interactive = command is None and entrypoint_override is None
+    if interactive and not inferred_interactive:
+        logging.error(
+            "Interactive SSH mode cannot be combined with --command or --entrypoint."
+        )
+        raise typer.Exit(code=1)
+    resolved_interactive = (
+        interactive if interactive is not None else inferred_interactive
+    )
+    noninteractive = not resolved_interactive
+    cmd_list: list[str] | None = shlex.split(command) if command is not None else None
+    ep_list: list[str] | None = (
+        shlex.split(entrypoint_override) if entrypoint_override is not None else None
+    )
+
+    public_key: str | None = None
+    if resolved_interactive:
+        if key:
+            if not key.exists():
+                logging.error(f"Public key not found: {key}")
+                raise typer.Exit(code=1)
+            public_key = key.read_text().strip()
+        else:
+            try:
+                public_key = FlowMesh().ssh.detect_public_key()
+            except FlowMeshError as exc:
+                logging.error(str(exc))
+                raise typer.Exit(code=1)
+
+        if mode not in ("direct", "proxy", "forward"):
+            logging.error(f"Invalid mode: {mode}. Use 'direct', 'proxy', or 'forward'.")
+            raise typer.Exit(code=1)
+
+    client = FlowMesh()
+    workflow_yaml = client.ssh.build_task_yaml(
+        name=name,
+        public_key=public_key,
+        user=user,
+        mode=mode,
+        ttl=ttl,
+        idle_timeout=idle_timeout,
+        gpu=gpu,
+        gpu_memory=gpu_memory,
+        cpu=cpu,
+        memory=memory,
+        image=image,
+        worker=worker,
+        env_pairs=env,
+        interactive=resolved_interactive,
+        command=cmd_list,
+        entrypoint=ep_list,
+    )
+
+    logging.info("Submitting SSH task...")
+    try:
+        result = client.workflows.submit(workflow_yaml)
+    except FlowMeshError as exc:
+        logging.error(f"Failed to submit SSH task: {exc}")
+        raise typer.Exit(code=1)
+
+    task_id = None
+    if result.tasks:
+        task_id = result.tasks[0].task_id
+    if not task_id:
+        logging.error("No task ID returned from submission.")
+        logging.error(result.model_dump_json(indent=2))
+        raise typer.Exit(code=1)
+
+    logging.info(f"Task submitted: {task_id}")
+
+    if noninteractive:
+        if logs:
+            _stream_logs_and_exit(client, task_id, interval, tail)
+        _poll_status_and_exit(client, task_id, interval)
+    else:
+        # Interactive: wait for SSH session and connect.
+        logging.info("Waiting for SSH session...")
+        try:
+            ssh_info = client.tasks.wait_for_ssh(task_id, interval=interval)
+        except FlowMeshError as exc:
+            logging.error(str(exc))
+            raise typer.Exit(code=1)
+        _exec_ssh(ssh_info, task_id, ssh_args, direct)
+
+
+@app.command("proxy")
+def proxy(
+    task_id: str = typer.Argument(..., help="Task identifier"),
+) -> None:
+    """Raw stdin/stdout proxy for SSH ProxyCommand (internal)."""
+    try:
+        asyncio.run(_run_proxy(task_id))
+    except KeyboardInterrupt:
+        pass
+
+
+async def _run_proxy(task_id: str) -> None:
+    """Async WebSocket <-> stdio proxy."""
+    client = FlowMesh()
+    api_key = client.api_key
+    ws_url = client.ssh.proxy_url(task_id)
+
+    auth_header = {"Authorization": f"Bearer {api_key}"} if api_key else None
+    async with websockets.connect(ws_url, additional_headers=auth_header) as ws:
+        loop = asyncio.get_running_loop()
+
+        async def stdin_to_ws() -> None:
+            try:
+                while True:
+                    data = await loop.run_in_executor(
+                        None, os.read, sys.stdin.fileno(), 4096
+                    )
+                    if not data:
+                        break
+                    await ws.send(data)
+            except Exception:
+                pass
+
+        async def ws_to_stdout() -> None:
+            try:
+                async for msg in ws:
+                    if isinstance(msg, str):
+                        msg = msg.encode()
+                    await loop.run_in_executor(None, os.write, sys.stdout.fileno(), msg)
+            except Exception:
+                pass
+
+        t1 = asyncio.create_task(stdin_to_ws())
+        t2 = asyncio.create_task(ws_to_stdout())
+        _, pending = await asyncio.wait([t1, t2], return_when=asyncio.FIRST_COMPLETED)
+        for task in pending:
+            task.cancel()
+            try:
+                await task
+            except (asyncio.CancelledError, Exception):
+                pass
+
+
+@app.command("connections")
+def list_connections(
+    connection_id: str | None = typer.Option(
+        None, "--id", help="Filter by connection id"
+    ),
+    access_mode: list[str] | None = typer.Option(
+        None, "--mode", help="Filter by access mode (repeatable)"
+    ),
+    task_id: str | None = typer.Option(None, "--task-id", help="Filter by task id"),
+    workflow_id: str | None = typer.Option(
+        None, "--workflow-id", help="Filter by workflow id"
+    ),
+    worker_id: str | None = typer.Option(
+        None, "--worker-id", help="Filter by worker id"
+    ),
+    node_id: str | None = typer.Option(None, "--node-id", help="Filter by node id"),
+    username: str | None = typer.Option(
+        None, "--username", help="Filter by SSH username"
+    ),
+    source_ip: str | None = typer.Option(
+        None, "--source-ip", help="Filter by source IP"
+    ),
+    source_port: int | None = typer.Option(
+        None, "--source-port", help="Filter by source port"
+    ),
+    query: list[str] | None = typer.Option(
+        None, "--query", "-q", help="Filter connections by key=value pairs"
+    ),
+) -> None:
+    """List active SSH connections audited by the server."""
+    client = FlowMesh()
+    query_params = parse_query_filters(query)
+    append_param(query_params, "connection_id", connection_id)
+    extend_params(query_params, "access_mode", access_mode)
+    append_param(query_params, "task_id", task_id)
+    append_param(query_params, "workflow_id", workflow_id)
+    append_param(query_params, "worker_id", worker_id)
+    append_param(query_params, "node_id", node_id)
+    append_param(query_params, "username", username)
+    append_param(query_params, "source_ip", source_ip)
+    append_param(query_params, "source_port", source_port)
+    try:
+        connections = client.ssh.list(query_params=query_params)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    logging.log(json.dumps([c.model_dump(mode="json") for c in connections], indent=2))

flowmesh_cli/commands/system.py

@@ -0,0 +1,22 @@
+import json
+
+import typer
+from flowmesh import FlowMesh
+from flowmesh.exceptions import FlowMeshError
+
+from ..core import logging
+from ..core.typer import get_typer
+
+app = get_typer(help="Query FlowMesh server system information.")
+
+
+@app.command()
+def metrics() -> None:
+    """Retrieve and display system metrics from the FlowMesh server."""
+    client = FlowMesh()
+    try:
+        result = client.system.metrics()
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    logging.log(json.dumps(result, indent=2))

flowmesh_cli/commands/task.py

@@ -0,0 +1,251 @@
+"""Public FlowMesh task commands."""
+
+import json
+import time
+from pathlib import Path
+
+import typer
+from flowmesh import FlowMesh
+from flowmesh.exceptions import FlowMeshError
+from flowmesh.models.common import TERMINAL_TASK_STATUSES, TaskStatus
+
+from ..core import logging
+from ..core.query import parse_query_filters
+from ..core.task import wait_for_task_completion
+from ..core.typer import get_typer
+
+app = get_typer(help="Query and monitor tasks executing on FlowMesh workers.")
+
+
+def _format_log_event(event: dict) -> str:
+    ts = event.get("ts", "")
+    message = str(event.get("message", "")).rstrip("\n")
+    prefix = f"[{ts}] " if ts else ""
+    return prefix + message
+
+
+def _log_ssh_connection_instructions(
+    task_id: str, latest_update: dict, client: FlowMesh
+) -> None:
+    ssh_info = latest_update.get("ssh")
+    if not isinstance(ssh_info, dict):
+        return
+    mode = str(ssh_info.get("mode") or "direct")
+    logging.log(f"[{task_id}] SSH connection instructions ({mode}):")
+    for label, command in client.ssh.connection_commands(task_id, ssh_info):
+        logging.log(f"  {label}:")
+        logging.log(f"    {command}")
+    if str(ssh_info.get("mode") or "direct") == "proxy":
+        logging.log("  note:")
+        logging.log("    proxy ssh command requires `websocat` and $FLOWMESH_API_KEY")
+
+
+@app.command()
+def info(task_id: str = typer.Argument(..., help="Task identifier")) -> None:
+    """Retrieve the current status and metadata for a specific task."""
+    client = FlowMesh()
+    try:
+        task = client.tasks.retrieve(task_id)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    logging.log(task.model_dump_json(indent=2))
+
+
+@app.command("list")
+def list_tasks(
+    task_id: str | None = typer.Option(None, "--id", help="Filter by task id"),
+    workflow_id: str | None = typer.Option(
+        None, "--workflow-id", help="Filter by workflow id"
+    ),
+    status: list[str] | None = typer.Option(
+        None, "--status", "-s", help="Filter by status (repeatable)"
+    ),
+    category: str | None = typer.Option(None, "--category", help="Filter by category"),
+    task_type: str | None = typer.Option(None, "--type", help="Filter by task type"),
+    assigned_worker: str | None = typer.Option(
+        None, "--assigned-worker", help="Filter by assigned worker id"
+    ),
+    graph_node_name: str | None = typer.Option(
+        None, "--graph-node", help="Filter by graph node name"
+    ),
+    completed: bool | None = typer.Option(
+        None,
+        "--completed/--not-completed",
+        help="Filter by completion state",
+    ),
+    failed: bool | None = typer.Option(
+        None,
+        "--failed/--not-failed",
+        help="Filter by failure state",
+    ),
+    query: list[str] | None = typer.Option(
+        None, "--query", "-q", help="Filter tasks by key=value pairs"
+    ),
+) -> None:
+    """List all tasks registered in the FlowMesh server."""
+    client = FlowMesh()
+    query_params = parse_query_filters(query)
+    try:
+        tasks = client.tasks.list(
+            task_id=task_id,
+            workflow_id=workflow_id,
+            status=status or None,
+            category=category,
+            task_type=task_type,
+            assigned_worker=assigned_worker,
+            graph_node_name=graph_node_name,
+            completed=completed,
+            failed=failed,
+            query_params=query_params,
+        )
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    logging.log(json.dumps([t.model_dump(mode="json") for t in tasks], indent=2))
+
+
+@app.command()
+def stop(
+    task_id: str = typer.Argument(..., help="Task identifier"),
+    interval: float = typer.Option(2.0, help="Polling interval in seconds"),
+    no_wait: bool = typer.Option(
+        False, "--no-wait", help="Return immediately after requesting stop"
+    ),
+) -> None:
+    """Stop a running task."""
+    client = FlowMesh()
+    try:
+        client.tasks.stop(task_id)
+    except FlowMeshError as exc:
+        logging.error(f"Failed to stop task: {exc}")
+        raise typer.Exit(code=1)
+
+    message = f"Stop requested for task {task_id}."
+    if no_wait:
+        logging.log(message)
+        return
+
+    logging.info(message)
+    status, error = wait_for_task_completion(task_id, interval)
+    match status:
+        case TaskStatus.DONE:
+            logging.success("Task stopped successfully.")
+        case TaskStatus.FAILED | TaskStatus.CANCELLED:
+            logging.error(f"Task {status.lower()}: {error}")
+            raise typer.Exit(code=1)
+        case _:
+            logging.error(f"Unexpected task status: {status}")
+            raise typer.Exit(code=1)
+
+
+@app.command()
+def watch(
+    task_id: str = typer.Argument(..., help="Task identifier"),
+    interval: float = typer.Option(2.0, help="Polling interval in seconds"),
+) -> None:
+    """Monitor a task's status by polling until completion."""
+    client = FlowMesh()
+    last_status: TaskStatus | None = None
+    last_update: dict | None = None
+    try:
+        while True:
+            task = client.tasks.retrieve(task_id)
+            payload = task.model_dump(mode="json")
+            status_value = task.status
+            if status_value != last_status:
+                logging.log(f"[{task_id}] status: {status_value}")
+                last_status = status_value
+            latest_update = payload.get("latest_update")
+            if isinstance(latest_update, dict) and latest_update != last_update:
+                logging.log(json.dumps(latest_update, indent=2))
+                _log_ssh_connection_instructions(task_id, latest_update, client)
+                last_update = latest_update
+            if status_value in TERMINAL_TASK_STATUSES:
+                logging.log(json.dumps(payload, indent=2))
+                return
+            time.sleep(interval)
+    except KeyboardInterrupt:
+        logging.warning("Cancelled by user.")
+        raise typer.Exit(code=1)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+
+
+logs_app = get_typer(help="Query and monitor task logs.")
+app.add_typer(logs_app, name="logs")
+
+
+@logs_app.command("show")
+def show_logs(
+    task_id: str = typer.Argument(..., help="Task identifier"),
+    limit: int = typer.Option(200, help="Maximum number of entries to return"),
+    before: str | None = typer.Option(None, help="Return entries before this cursor"),
+    after: str | None = typer.Option(None, help="Return entries after this cursor"),
+    json_output: bool = typer.Option(
+        False, "--json", help="Print raw JSON response instead of formatted lines"
+    ),
+) -> None:
+    """Query recent task logs from the server."""
+    client = FlowMesh()
+    try:
+        result = client.tasks.get_logs(task_id, limit=limit, before=before, after=after)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    payload = result.model_dump(mode="json")
+    if json_output:
+        logging.log(json.dumps(payload, indent=2))
+        return
+    entries = payload.get("entries") or []
+    for entry in entries:
+        if not isinstance(entry, dict):
+            continue
+        event = entry.get("event")
+        if not isinstance(event, dict):
+            continue
+        logging.log(_format_log_event(event))
+    next_cursor = payload.get("next_cursor")
+    prev_cursor = payload.get("prev_cursor")
+    if next_cursor or prev_cursor:
+        logging.log(
+            json.dumps({"next_cursor": next_cursor, "prev_cursor": prev_cursor})
+        )
+
+
+@logs_app.command("stream")
+def stream_logs(
+    task_id: str = typer.Argument(..., help="Task identifier"),
+    cursor: str | None = typer.Option(None, help="Start streaming after this cursor"),
+) -> None:
+    """Stream task logs via SSE."""
+    client = FlowMesh()
+    try:
+        for entry in client.tasks.stream_logs(task_id, cursor=cursor):
+            event = entry.event.model_dump(mode="json")
+            logging.log(_format_log_event(event))
+    except KeyboardInterrupt:
+        logging.warning("Cancelled by user.")
+        raise typer.Exit(code=1)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+
+
+@logs_app.command("download")
+def download_logs(
+    task_id: str = typer.Argument(..., help="Task identifier"),
+    output: Path = typer.Option(..., "--output", "-o", help="Path to save logs.jsonl"),
+) -> None:
+    """Download archived logs.jsonl for a task."""
+    client = FlowMesh()
+    try:
+        client.tasks.download_logs(task_id, output)
+    except FlowMeshError as exc:
+        logging.error(str(exc))
+        raise typer.Exit(code=1)
+    except OSError as exc:
+        logging.error(f"Failed to write {output}: {exc}")
+        raise typer.Exit(code=1)
+    logging.log(f"Wrote logs to {output}")