tuneplane-node 0.3.15__py3-none-any.whl

tuneplane_node/cli.py

@@ -0,0 +1,189 @@
+"""tuneplane-node: the entry point of the daemon that runs on a GPU node.
+
+    TUNEPLANE_NODE_TOKEN=<shared token> tuneplane-node serve --host 0.0.0.0 --port 7070
+
+The node side takes the same TUNEPLANE_LOCAL_* settings as the local backend (container
+runtime, card count, health checks), plus TUNEPLANE_STORAGE_ROOT (shared storage, mounted
+at the same path as on the console).
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import typer
+
+app = typer.Typer(add_completion=False, no_args_is_help=True, help="tuneplane agent daemon")
+
+
+@app.callback()
+def _main() -> None:
+    """Keep `serve` a subcommand.
+
+    Typer folds a lone command into the root, which made the documented
+    `tuneplane-node serve --port 7070` fail with "Got unexpected extra argument
+    (serve)" -- in the published operator docs, this module's own docstring, and
+    any systemd unit written from either. server/cli.py carries the same guard for
+    the same reason; the agent was simply missing it.
+    """
+
+
+def _state_dir(settings) -> Path:
+    return Path(str(settings.node_state_dir)).expanduser()
+
+
+@app.command()
+def serve(
+    host: str = typer.Option("0.0.0.0", help="Address to bind"),
+    port: int = typer.Option(7070, help="Port to listen on"),
+) -> None:
+    """Start the agent daemon. Configured by env / .env, the same way the console is.
+
+    If this node has joined a fleet, a heartbeat runs alongside the daemon. The
+    daemon serves whether or not it has: losing the console is not a reason to
+    stop running work, and the containers on this machine stay the truth about
+    what is running (ADR-002).
+    """
+    import asyncio
+
+    import uvicorn
+
+    from tuneplane_node.daemon import create_node_app
+    from tuneplane_node.join import heartbeat_forever, load_identity
+    from tuneplane_node.settings import NodeSettings
+
+    settings = NodeSettings()
+    api = create_node_app(settings)
+    identity = load_identity(_state_dir(settings))
+
+    if identity is not None:
+        address = f"http://{host}:{port}" if host not in ("0.0.0.0", "::") else ""
+
+        @api.on_event("startup")
+        async def _start_heartbeat() -> None:  # pragma: no cover - process wiring
+            api.state.heartbeat = asyncio.create_task(
+                heartbeat_forever(
+                    identity, settings, node_version=_node_version(), address=address
+                )
+            )
+
+        @api.on_event("shutdown")
+        async def _stop_heartbeat() -> None:  # pragma: no cover - process wiring
+            task = getattr(api.state, "heartbeat", None)
+            if task is not None:
+                task.cancel()
+
+        typer.echo(f"joined fleet {identity.fleet_id} as node {identity.node_id}")
+    else:
+        typer.echo("not joined to any fleet; run `tuneplane-node join` to register this machine")
+
+    uvicorn.run(api, host=host, port=port, log_level="info")
+
+
+@app.command()
+def join(
+    server: str = typer.Option(
+        ..., "--server",
+        help="Where the control plane answers, e.g. https://tuneplane.corp",
+    ),
+    fleet: str = typer.Option(..., "--fleet", help="Fleet to join"),
+    token: str = typer.Option(..., "--token", help="Join token an administrator minted"),
+    address: str = typer.Option(
+        "", "--address",
+        help="Where the console reaches this machine, e.g. http://gpu-1:7070. Defaults to this "
+             "host's name and the agent port.",
+    ),
+    port: int = typer.Option(7070, "--port", help="Port the agent daemon listens on"),
+    force: bool = typer.Option(False, "--force", help="Re-join even though this machine already has an identity"),
+) -> None:
+    """Register this machine into a fleet.
+
+    **No platform login.** The join token an administrator gave you is the whole
+    authorization -- you need root on this box, not an account on the console
+    (ADR-0013).
+
+    Run it once per machine. What comes back is this node's identity and a
+    credential, written to the agent state directory; `tuneplane-node serve` then
+    heartbeats with them.
+    """
+    import asyncio
+    import socket
+
+    from tuneplane_node.inventory import collect
+    from tuneplane_node.join import JoinError, load_identity, save_identity
+    from tuneplane_node.join import join as do_join
+    from tuneplane_node.settings import NodeSettings
+
+    settings = NodeSettings()
+    state_dir = _state_dir(settings)
+
+    existing = load_identity(state_dir)
+    if existing is not None and not force:
+        # Re-joining silently would strand the old identity: the console still
+        # believes that node exists, and its containers would reconcile against
+        # a node id nothing reports on any more.
+        typer.echo(
+            f"this machine is already node {existing.node_id} in fleet {existing.fleet_id}. "
+            f"Pass --force to register it again as a new node.",
+            err=True,
+        )
+        raise typer.Exit(2)
+
+    reachable = address or f"http://{socket.gethostname()}:{port}"
+
+    async def _run():
+        inventory = await collect(settings)
+        return inventory, await do_join(
+            server,
+            fleet_id=fleet,
+            token=token,
+            address=reachable,
+            inventory=inventory.to_dict(),
+            node_version=_node_version(),
+        )
+
+    try:
+        inventory, (identity, warnings) = asyncio.run(_run())
+    except JoinError as exc:
+        typer.echo(str(exc), err=True)
+        raise typer.Exit(2) from exc
+
+    path = save_identity(state_dir, identity)
+    typer.echo(f"joined fleet {fleet} as node {identity.node_id}")
+    cards = inventory.accelerator_name or inventory.accelerator_series or "unknown"
+    typer.echo(f"  {inventory.accelerator_count} x {cards}")
+    if inventory.accelerator_name and not inventory.accelerator_series:
+        # The console can place a job on this node, but not one that pins a
+        # series -- and the operator is the only person who can fix that.
+        typer.echo(
+            f"  ! this card is not in the platform's hardware registry, so jobs that pin a "
+            f"series will not be placed here. Ask an administrator to add {cards!r}."
+        )
+    typer.echo(f"  driver {inventory.driver or 'unknown'}, runtime {inventory.runtime or 'unknown'}")
+    for warning in warnings:
+        typer.echo(f"  ! {warning}", err=True)
+    typer.echo(
+        f"  storage root {'visible' if inventory.storage_root_visible else 'NOT visible'}"
+    )
+    typer.echo(f"  identity written to {path}")
+    typer.echo("  start the agent with: tuneplane-node serve")
+
+
+def _node_version() -> str:
+    """This the node daemon's own version, reported to the console on every heartbeat.
+
+    Its own distribution, never the console's: a machine that installed
+    `tuneplane-node` has no `server` package to ask, and reaching for one
+    here would be the exact dependency direction this package exists to remove.
+    """
+    from importlib.metadata import PackageNotFoundError, version
+
+    try:
+        return version("tuneplane-node")
+    except PackageNotFoundError:
+        # Running from a checkout. Empty rather than a guess: the console shows
+        # it verbatim, and an invented version is worse than a blank.
+        return ""
+
+
+if __name__ == "__main__":
+    app()