app/backend/fluksio/cli.py

"""`fluksio serve`, `fluksio enroll`, `fluksio worker`, `sync`, `run`, `runs`.

The point of this module is a machine nobody can route to: a node on a cluster
where ports cannot be opened, or a laptop with no Docker. `fluksio serve`
starts the engine with no infrastructure and no configuration; `fluksio enroll`
hands it a claim code, and it dials the portal itself. What a browser would
have done locally is then done through the portal, which serves the dashboard
from its own side.

Nothing from the engine is imported at module level. `fluksio.core.config`
builds its settings when it is first imported, and the database engine and the
user venv's path are computed from those — so the environment has to be right
before any of that happens, which is what `_configure_environment` is for.
"""

from __future__ import annotations

import argparse
import copy
import json
import os
import secrets
import socket
import sys
from collections.abc import Callable
from pathlib import Path
from typing import Any

import fluksio

#: Where an installation keeps everything, unless it is told otherwise.
DEFAULT_HOME = Path("~/.fluksio")

#: The portal a claim code is redeemed at, unless another is named. Having a
#: default is the difference between one flag and two on the one command that
#: is run before anything works.
DEFAULT_PORTAL = "https://hub.fluksio.com"

#: WAL — what lets readers work while the engine writes — is not supported on
#: these. The database would be corrupt or locked, so it is worth saying.
NETWORK_FILESYSTEMS = ("nfs", "nfs4", "cifs", "smb", "smb3", "lustre", "fuse.sshfs")


def _say(message: str = "") -> None:
    """Print, and mean it.

    Redirected to a file or a journal, stdout is block-buffered, and the admin
    password below is shown exactly once — sitting in a buffer until the
    process exits is the same as never printing it.
    """
    print(message, flush=True)


def _data_dir(raw: str | None, shared: bool = False) -> Path:
    """Which installation this command is for.

    A repository with its own venv wants its own engine too — its own flows,
    its own run history, its own token — so the default is a `.fluksio` beside
    the code, found the way `.git` is. `--global` asks for the shared one
    instead, and `--data-dir` names any directory outright.
    """
    from fluksio.sdk.client import find_data_dir, ignore_self

    if raw:
        path = Path(raw).expanduser()
    elif shared:
        path = DEFAULT_HOME.expanduser()
    else:
        found = find_data_dir()
        path = found if found is not None else Path.cwd() / ".fluksio"
    path.mkdir(parents=True, exist_ok=True)
    ignore_self(path)
    return path.resolve()


def _mention_other_installation(data_dir: Path) -> None:
    """Say when there is a second engine, so neither goes looking lost."""
    shared = DEFAULT_HOME.expanduser().resolve()
    if data_dir == shared or not (shared / "fluksio.db").exists():
        return
    _say(f"  Note     {shared} holds another installation; this one is separate.")
    _say("           `fluksio serve --global` runs that one instead.")


def _mention_undeclared_cards() -> None:
    """Say when a stored flow asks for a card this engine does not have.

    Cards are declared rather than detected, so an engine told nothing has
    none — and a node asking for one is clamped to zero and runs beside every
    other, which on a GPU is the deadlock the declaration exists to prevent.
    The placer says so once it happens, into the log; this says it while
    somebody is still reading the terminal.
    """
    from fluksio.core.config import settings

    if settings.FLOW_GPUS:
        return
    from fluksio.flow.runs import required_resources
    from fluksio.flow.store import FlowStore

    store = FlowStore(settings.FLOWS_DIR)
    asking = []
    for name in store.list_flows():
        try:
            needs = required_resources(store.read_flow(name))
        except Exception:
            # A flow that will not parse is the engine's to complain about.
            continue
        if (needs or {}).get("gpus"):
            asking.append(name)
    if not asking:
        return
    named = ", ".join(sorted(asking)[:3]) + (" …" if len(asking) > 3 else "")
    _say(f"  Cards    0 declared, but {named} asks for one.")
    _say("           Nothing detects them: `--gpus N` says how many are here.")


def _warn_if_networked(path: Path) -> None:
    """A cluster's $HOME is often NFS, and SQLite's WAL does not work there."""
    try:
        mounts = Path("/proc/mounts").read_text().splitlines()
    except OSError:  # pragma: no cover - not Linux
        return
    best, kind = "", ""
    for line in mounts:
        parts = line.split()
        if len(parts) < 3:
            continue
        point, fstype = parts[1], parts[2]
        if (path == Path(point) or point in map(str, path.parents)) and len(
            point
        ) > len(best):
            best, kind = point, fstype
    if kind in NETWORK_FILESYSTEMS:
        print(
            f"warning: {path} is on {kind}, where SQLite's write-ahead log does "
            "not work. Point --data-dir at local disk.",
            file=sys.stderr,
            flush=True,
        )


def load_or_create_secret_key(path: Path) -> str:
    """The key that signs sessions and encrypts the secrets store, kept once.

    Regenerating it per process would sign out every session on restart and,
    worse, leave `secrets.enc` unreadable.
    """
    if path.exists():
        return path.read_text().strip()
    key = secrets.token_urlsafe(32)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(key)
    path.chmod(0o600)
    return key


def _configure_environment(data_dir: Path) -> None:
    """Everything the settings need, before anything reads them.

    Imports nothing from the engine, deliberately: the first import of
    `fluksio.core.config` builds the settings, and any engine module reaches it
    within an import or two. Doing that here would fix `DATA_DIR` at whatever
    the current directory happened to be.

    `setdefault` throughout: an operator who exported one of these means it.
    """
    os.environ.setdefault("DATA_DIR", str(data_dir))
    # Without this the settings would read whatever `../.env` resolves to from
    # the current directory — a checkout's development configuration, if the
    # command happened to be run from inside one.
    os.environ.setdefault("FLUKSIO_ENV_FILE", str(data_dir / "env"))
    os.environ.setdefault(
        "SECRET_KEY", load_or_create_secret_key(data_dir / "secret_key")
    )


def _prepare(data_dir: Path) -> None:
    _warn_if_networked(data_dir)
    _configure_environment(data_dir)
    from fluksio.core.db import engine, prepare

    prepare(engine)


def _enroll(portal: str, code: str, as_email: str | None) -> int:
    from sqlmodel import Session

    from fluksio.cloud import enroll as enroll_mod
    from fluksio.core.bootstrap import ensure_superuser, pick_superuser
    from fluksio.core.config import settings
    from fluksio.core.db import engine

    with Session(engine) as session:
        admin, generated = ensure_superuser(session, email=as_email)
        if generated:
            _print_new_admin(admin.email, generated)
        try:
            user = pick_superuser(session, as_email)
        except LookupError as exc:
            print(f"error: {exc}", file=sys.stderr, flush=True)
            return 1
        # Read before the session closes: the instance is detached after it,
        # and touching an attribute then goes back to a database that is gone.
        email = user.email
        try:
            config = enroll_mod.enroll(session, user, portal, code)
        except enroll_mod.AlreadyEnrolled as exc:
            print(
                f"error: {exc.detail} (see {settings.CLOUD_CONFIG_FILE}).",
                file=sys.stderr,
                flush=True,
            )
            return 1
        except enroll_mod.EnrollError as exc:
            print(f"error: {exc.detail}", file=sys.stderr, flush=True)
            return 1

    _say(
        f"Connected to {config.portal_url} as {email} "
        f"(installation {config.installation_id})."
    )
    return 0


def _print_new_admin(email: str, password: str) -> None:
    _say(f"Created the admin account {email}")
    _say(f"  password: {password}")
    _say("  Shown once. Change it from the dashboard.")


#: How long the token `serve` writes for its own machine is good for. Long,
#: because the thing it saves you from is logging in again, and it is rewritten
#: on every start anyway — an engine left running for a year is the only case
#: this has to cover.
LOCAL_TOKEN_DAYS = 365


def _sign_in(admin_id: Any, url: str, data_dir: Path) -> Path:
    """Write the credential for the engine this command is about to start.

    Logging in to your own machine is a formality: the password was printed by
    this same process, and the database it authenticates against is in the
    directory the token goes into. So it costs nothing to be honest about it
    and hand the client a token, rather than asking somebody to type back
    something we just told them.
    """
    from datetime import timedelta

    from fluksio.core import security
    from fluksio.sdk.client import write_config

    token = security.create_access_token(
        admin_id, expires_delta=timedelta(days=LOCAL_TOKEN_DAYS)
    )
    return write_config(url, token, data_dir)


#: The `serve` flags that are settings under another name, and the setting each
#: one writes. A flag is a real environment variable, which outranks the env
#: file the settings read — so the order is flag, environment, `<data-dir>/env`.
CONCURRENCY_FLAGS = {
    "max_workers": "FLOW_MAX_WORKERS",
    "max_cascades": "FLOW_MAX_CASCADES",
    "max_runs": "FLOW_MAX_RUNS",
    "gpus": "FLOW_GPUS",
}


def _at_least(minimum: int) -> Callable[[str], int]:
    """A flag's value, checked here rather than by the settings.

    These are written into the environment before the settings are built, so a
    number they refuse dies inside a pydantic import with no flag named in it.
    """

    def parse(text: str) -> int:
        value = int(text)
        if value < minimum:
            raise argparse.ArgumentTypeError(f"is {value}, needs at least {minimum}")
        return value

    return parse


#: What `serve` listens on when nobody says. Taken often enough — another
#: engine, another framework's dev server — that dying on it is the first
#: thing a zero-config start would hit.
DEFAULT_PORT = 8000

#: How far up from it to look before giving up and letting the bind fail.
PORT_TRIES = 20


#: Where a serving engine records itself, beside the data it is serving. Read
#: to tell "this directory's engine is already up" from "something else has the
#: port", which are the two ways a second `serve` fails to be what was wanted.
PIDFILE = "serve.pid"


def write_pidfile(data_dir: Path, port: int) -> Path:
    """Record which process is serving this directory, and where."""
    path = data_dir / PIDFILE
    path.write_text(json.dumps({"pid": os.getpid(), "port": port}))
    return path


def read_pidfile(data_dir: Path) -> dict[str, int] | None:
    """The engine serving this directory, if one still is.

    A process killed outright leaves the file behind, so the pid is checked
    rather than believed — a stale file is the same as no file.
    """
    try:
        record = json.loads((data_dir / PIDFILE).read_text())
        pid, port = int(record["pid"]), int(record["port"])
    except (OSError, ValueError, KeyError, TypeError):
        return None
    try:
        os.kill(pid, 0)
    except (OSError, ProcessLookupError):
        return None
    return {"pid": pid, "port": port}


def _token_for(data_dir: Path) -> str:
    """The credential this directory's last engine wrote, if it wrote one."""
    from fluksio.sdk.client import config_path

    try:
        return str(json.loads(config_path(data_dir).read_text()).get("token", ""))
    except (OSError, ValueError, AttributeError):
        return ""


def probe_engine(url: str, token: str, client: Any = None) -> str:
    """Who is on this port: ``ours``, ``foreign``, or ``other``.

    ``ours`` means an engine serving *this* data directory, which is what
    makes stopping it something this command may offer. The proof is the
    token: it is signed with this directory's secret key, so an engine that
    accepts it is one reading this directory's database. A Fluksio belonging
    to another installation answers the health check and refuses the token,
    and is only ever named — never stopped from here.
    """
    import logging

    import httpx

    # Two requests nobody asked for, on a start that is otherwise quiet until
    # uvicorn's own banner. httpx logs every one of them at INFO.
    noisy = logging.getLogger("httpx")
    was = noisy.level
    noisy.setLevel(logging.WARNING)
    http = client or httpx.Client(timeout=2.0)
    try:
        health = http.get(f"{url}/api/v1/utils/health-check/")
        if health.status_code != 200:
            return "other"
        # No token at all is asked unauthenticated: `Bearer ` is not a legal
        # header value, and a Fluksio this directory cannot prove is its own
        # is foreign — which is the answer that never stops anything.
        auth = {"Authorization": f"Bearer {token}"} if token else {}
        answer = http.get(f"{url}/api/v1/observability/summary", headers=auth)
        return "ours" if answer.status_code == 200 else "foreign"
    except Exception:
        return "other"
    finally:
        noisy.setLevel(was)
        if client is None:
            http.close()


def _free_port(host: str, start: int) -> int:
    """The first port from ``start`` that nothing is listening on.

    Probed with the same address and options uvicorn will bind with, so this
    answers the question uvicorn is about to ask rather than a similar one.
    """
    for port in range(start, start + PORT_TRIES):
        with socket.socket() as probe:
            probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            try:
                probe.bind((host, port))
            except OSError:
                continue
            return port
    return start


def cmd_serve(args: argparse.Namespace) -> int:
    # At a terminal this is a dashboard with the engine as a child of it. The
    # import is here rather than at the top because it is only ever needed on
    # that path, and `serve` in a container must not pay for it.
    if not args.plain and sys.stdout.isatty() and sys.stdin.isatty():
        from fluksio.tui import run_tui

        return run_tui(args)

    data_dir = _data_dir(args.data_dir, args.shared)
    for flag, name in CONCURRENCY_FLAGS.items():
        value = getattr(args, flag, None)
        if value is not None:
            os.environ[name] = str(value)
    _prepare(data_dir)

    from sqlmodel import Session

    from fluksio.cloud import config as cloud_config
    from fluksio.core.bootstrap import ensure_superuser
    from fluksio.core.db import engine

    with Session(engine) as session:
        admin, generated = ensure_superuser(
            session, email=args.admin_email, password=args.admin_password
        )
        admin_id, admin_email = admin.id, admin.email
        if generated:
            _print_new_admin(admin_email, generated)

    if args.enroll:
        if not cloud_config.exists():
            # Before the engine starts, so the connector finds the config and
            # dials out as part of coming up rather than needing a restart.
            failed = _enroll(
                args.portal or DEFAULT_PORTAL, args.enroll, args.admin_email
            )
            if failed:
                return failed

    import uvicorn

    from fluksio.flow import modules
    from fluksio.main import app

    # The client talks to this engine, and 0.0.0.0 is not an address to talk
    # to — it is a statement about which interfaces to listen on.
    reachable = "127.0.0.1" if args.host in ("0.0.0.0", "::", "") else args.host

    port = args.port
    if port is None:
        port = _free_port(args.host, DEFAULT_PORT)
        if port != DEFAULT_PORT:
            # Moving off the port quietly makes starting a second engine for
            # one directory look like it worked. Two of them on one SQLite
            # file is not a supported shape — two *installations* on one
            # machine is — so the one already up is named instead.
            url = f"http://{reachable}:{DEFAULT_PORT}"
            who = probe_engine(url, _token_for(data_dir))
            if who == "ours":
                running = read_pidfile(data_dir)
                where = f" (pid {running['pid']})" if running else ""
                _say(f"An engine for {data_dir} is already serving at {url}{where}.")
                _say("  fluksio status talks to it; stop it to start another.")
                return 0
            if who == "foreign":
                _say(
                    f"Port {DEFAULT_PORT} holds another installation's Fluksio; "
                    f"serving on {port} instead."
                )
            else:
                _say(f"Port {DEFAULT_PORT} is in use; serving on {port} instead.")

    url = f"http://{reachable}:{port}"
    token_path = _sign_in(admin_id, url, data_dir)

    config = cloud_config.load()
    _say(f"Fluksio {fluksio.__version__} — data in {data_dir}")
    _say(f"  API      {url}/api/v1")
    # Which environment node code imports from is the thing a data scientist
    # most needs to know at this moment, and the answer differs depending on
    # how Fluksio was installed. Saying it costs one line.
    if modules.adopted() is not None:
        _say(f"  Nodes    {modules.venv_python()}")
        _say("           your environment, adopted. Add packages with pip.")
    else:
        # Not `venv_python()`: the managed venv is built when the engine comes
        # up, which is after this prints, and until then that would answer with
        # whatever interpreter happens to be running this.
        _say(f"  Nodes    {modules.venv_dir() / 'bin' / 'python'}")
        _say("           a venv of its own; the Modules screen installs into it.")
    if config is not None:
        _say(f"  Portal   {config.portal_url}, installation {config.installation_id}")
        _say("  The dashboard is served by the portal; nothing is served here.")
    else:
        _say("  No portal. Pair this installation with:")
        _say("    fluksio enroll <code>")
    _say(f"  Signed in as {admin_email}")
    _say(f"           token in {token_path}")
    _mention_undeclared_cards()
    _mention_other_installation(data_dir)

    # One process: it holds the flow engine, and a second worker would be a
    # second engine — duplicated subscriptions, cron ticks and webhooks.
    pidfile = write_pidfile(data_dir, port)
    try:
        uvicorn.run(
            app,
            host=args.host,
            port=port,
            log_level=args.log_level,
            log_config=_log_config(args.log_level),
        )
    finally:
        pidfile.unlink(missing_ok=True)
    return 0


def _log_config(level: str) -> dict[str, Any]:
    """Uvicorn's logging, with the engine's own loggers drawn the same way.

    Without this the engine's lines go to the root logger, which has no handler
    configured and falls back to `INFO:fluksio.cloud.connector:...` — beside
    uvicorn's own aligned, coloured output it reads like something went wrong.
    """
    from uvicorn.config import LOGGING_CONFIG

    config = copy.deepcopy(LOGGING_CONFIG)
    # Named rather than configuring the root: everything else that logs — httpx
    # on every portal call, for one — is at INFO too, and today none of it is
    # printed at all. Styling the root would turn all of it on.
    for name in ("fluksio", "alembic"):
        config["loggers"][name] = {
            "handlers": ["default"],
            "level": level.upper(),
            # It has a handler of its own now; propagating would print each
            # line twice the moment anything configures the root.
            "propagate": False,
        }
    return config


def cmd_enroll(args: argparse.Namespace) -> int:
    data_dir = _data_dir(args.data_dir, args.shared)
    _prepare(data_dir)
    result = _enroll(args.portal, args.code, args.as_email)
    if result == 0:
        # True either way round: one already serving notices within seconds,
        # and one not yet started dials as it comes up.
        _say("An engine already running picks this up; otherwise `fluksio serve`.")
    return result


def cmd_worker(rest: list[str]) -> int:
    from fluksio_worker.agent import main as worker_main

    return worker_main(rest)


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="fluksio", description="Node-based automation: flows, dashboards, runs."
    )
    parser.add_argument("--version", action="version", version=fluksio.__version__)
    subparsers = parser.add_subparsers(dest="command", required=True)

    def with_data_dir(sub: argparse.ArgumentParser) -> None:
        sub.add_argument(
            "--data-dir",
            default=os.environ.get("FLUKSIO_HOME"),
            help="where this installation keeps everything (default: ./.fluksio)",
        )
        sub.add_argument(
            "--global",
            dest="shared",
            action="store_true",
            help=f"use the shared installation at {DEFAULT_HOME} instead",
        )

    serve = subparsers.add_parser("serve", help="run the engine")
    with_data_dir(serve)
    serve.add_argument("--host", default="127.0.0.1")
    # No default: a port nobody asked for may move when it is taken, and one
    # that was asked for may not.
    serve.add_argument(
        "--port",
        type=int,
        default=None,
        help=f"default {DEFAULT_PORT}, or the next free port when it is in use",
    )
    serve.add_argument("--log-level", default="info")
    serve.add_argument("--admin-email", default=None)
    serve.add_argument("--admin-password", default=None)
    serve.add_argument("--enroll", metavar="CODE", help="claim code, if not yet paired")
    serve.add_argument(
        "--portal",
        metavar="URL",
        help=f"the portal --enroll redeems at (default {DEFAULT_PORTAL})",
    )
    serve.add_argument(
        "--max-runs",
        type=_at_least(1),
        default=None,
        metavar="N",
        help="batch runs driven at once (default 4, FLOW_MAX_RUNS)",
    )
    serve.add_argument(
        "--max-cascades",
        type=_at_least(1),
        default=None,
        metavar="N",
        help="cascades in flight at once (default 4, FLOW_MAX_CASCADES)",
    )
    serve.add_argument(
        "--max-workers",
        type=_at_least(1),
        default=None,
        metavar="N",
        help="python worker processes (default 4, FLOW_MAX_WORKERS)",
    )
    serve.add_argument(
        "--plain",
        action="store_true",
        help="the log stream rather than the dashboard (the default with no terminal)",
    )
    serve.add_argument(
        "--gpus",
        type=_at_least(0),
        default=None,
        metavar="N",
        help="GPUs on this machine a node may be given (default 0, FLOW_GPUS)",
    )
    serve.set_defaults(func=cmd_serve)

    enroll = subparsers.add_parser(
        "enroll", help="pair this installation with a portal"
    )
    enroll.add_argument("code", help="the claim code minted on the portal")
    enroll.add_argument(
        "--portal",
        default=DEFAULT_PORTAL,
        metavar="URL",
        help=f"a portal of your own, instead of {DEFAULT_PORTAL}",
    )
    enroll.add_argument(
        "--as",
        dest="as_email",
        default=None,
        help="the local account a portal session arrives as",
    )
    with_data_dir(enroll)
    enroll.set_defaults(func=cmd_enroll)

    subparsers.add_parser(
        "worker",
        help="run nodes for an engine elsewhere (fluksio-worker)",
        add_help=False,
    )

    # The client half: talking to an engine rather than being one.
    from fluksio.sdk.cli import add_parsers

    add_parsers(subparsers)
    return parser


def main(argv: list[str] | None = None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    # Everything after `worker` belongs to the agent's own parser.
    if argv and argv[0] == "worker":
        return cmd_worker(argv[1:])
    parser = _parser()
    if argv and argv[0] == "run":
        # A flow's inputs are its own, so `--lr 0.05` cannot be declared here:
        # whatever this parser does not know is typed against the flow.
        args, rest = parser.parse_known_args(argv)
        return int(args.func(args, rest))
    args = parser.parse_args(argv)
    result: int = args.func(args)
    return result


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