Spaces:
Running
Running
| """python -m msgboard {init,serve,restore,backup,rotate-operator-token,demo,deploy-space}""" | |
| import argparse | |
| import ipaddress | |
| import os | |
| import signal | |
| import sys | |
| import time | |
| from pathlib import Path | |
| from . import db as dbmod | |
| from . import store | |
| from .config import DB_NAME, DEFAULT_PORT, OPERATOR_TOKEN_FILE | |
| NETWORK_FS = {"nfs", "nfs4", "cifs", "smb3", "fuse", "fuseblk", "9p", "sshfs"} | |
| def _fs_type(path: Path) -> str | None: | |
| """Best-effort filesystem type of `path` from /proc/mounts (Linux only).""" | |
| try: | |
| mounts = Path("/proc/mounts").read_text().splitlines() | |
| except OSError: | |
| return None | |
| best, fstype = "", None | |
| target = str(path.resolve()) | |
| for line in mounts: | |
| parts = line.split() | |
| if len(parts) >= 3 and (target == parts[1] or target.startswith(parts[1].rstrip("/") + "/")): | |
| if len(parts[1]) > len(best): | |
| best, fstype = parts[1], parts[2] | |
| return fstype | |
| def _warn_fs(data_dir: Path) -> None: | |
| fstype = _fs_type(data_dir) | |
| if fstype and (fstype in NETWORK_FS or fstype.startswith("fuse")): | |
| print(f"warning: {data_dir} is on a {fstype} filesystem. SQLite WAL needs a local disk; " | |
| "do not keep a live board on network or bucket mounts.", file=sys.stderr) | |
| def _write_secret_file(path: Path, secret: str) -> None: | |
| """Create a 0600 file that must not already exist, and fsync it.""" | |
| fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) | |
| try: | |
| os.write(fd, (secret + "\n").encode()) | |
| os.fsync(fd) | |
| finally: | |
| os.close(fd) | |
| def cmd_init(args) -> int: | |
| data_dir = Path(args.data_dir).expanduser() | |
| data_dir.mkdir(parents=True, exist_ok=True) | |
| os.chmod(data_dir, 0o700) | |
| _warn_fs(data_dir) | |
| db_path = data_dir / DB_NAME | |
| token_path = data_dir / OPERATOR_TOKEN_FILE | |
| existed = db_path.exists() | |
| conn = dbmod.connect(db_path) | |
| try: | |
| before = dbmod.schema_version(conn) | |
| try: | |
| version = dbmod.migrate(conn, db_path) | |
| except dbmod.SchemaError as exc: | |
| print(f"error: {exc}", file=sys.stderr) | |
| return 1 | |
| if store.admin_exists(conn): | |
| note = "migrated" if before != version else "unchanged" | |
| print(f"Board already initialized in {data_dir} (schema v{version}, {note}).") | |
| tmp_path = data_dir / (OPERATOR_TOKEN_FILE + ".tmp") | |
| if not token_path.exists() and tmp_path.exists(): | |
| # An earlier init committed but stopped before moving the file into place. | |
| try: | |
| store.authenticate(conn, tmp_path.read_text().strip(), time.time()) | |
| os.link(tmp_path, token_path) | |
| tmp_path.unlink() | |
| except Exception: | |
| pass | |
| if token_path.exists(): | |
| print(f"Operator token file: {token_path}") | |
| else: | |
| print("Operator token file is missing; existing credentials were left untouched. " | |
| "Run `python -m msgboard rotate-operator-token --data-dir ...` to issue a new one.") | |
| return 0 | |
| if token_path.exists(): | |
| print(f"error: {token_path} exists but the database has no operator; refusing to overwrite it. " | |
| "Move the file away if it is stale.", file=sys.stderr) | |
| return 1 | |
| # Write the secret to a temp file first, then commit, then move it into place. | |
| tmp_path = data_dir / (OPERATOR_TOKEN_FILE + ".tmp") | |
| if tmp_path.exists(): | |
| tmp_path.unlink() | |
| conn.execute("BEGIN IMMEDIATE") | |
| try: | |
| _, token = store.create_operator(conn, time.time()) | |
| _write_secret_file(tmp_path, token) | |
| conn.execute("COMMIT") | |
| except BaseException: | |
| conn.execute("ROLLBACK") | |
| if tmp_path.exists(): | |
| tmp_path.unlink() | |
| raise | |
| os.link(tmp_path, token_path) # fails rather than overwrite | |
| tmp_path.unlink() | |
| finally: | |
| conn.close() | |
| print(f"{'Initialized schema in existing' if existed else 'Created'} board in {data_dir} (schema v{version}).") | |
| print(f"Operator token saved to {token_path} (mode 0600). Keep it secret; open the dashboard with it to create invite links.") | |
| return 0 | |
| def cmd_rotate_operator(args) -> int: | |
| data_dir = Path(args.data_dir).expanduser() | |
| db_path = data_dir / DB_NAME | |
| token_path = data_dir / OPERATOR_TOKEN_FILE | |
| if not db_path.exists(): | |
| print(f"error: no board in {data_dir}", file=sys.stderr) | |
| return 1 | |
| conn = dbmod.connect(db_path) | |
| try: | |
| if dbmod.schema_version(conn) != dbmod.SCHEMA_VERSION: | |
| print("error: run init first to migrate the schema", file=sys.stderr) | |
| return 1 | |
| new_path = data_dir / (OPERATOR_TOKEN_FILE + ".new") | |
| if new_path.exists(): | |
| new_path.unlink() | |
| conn.execute("BEGIN IMMEDIATE") | |
| try: | |
| token = store.rotate_operator(conn, time.time()) | |
| _write_secret_file(new_path, token) | |
| conn.execute("COMMIT") | |
| except BaseException: | |
| conn.execute("ROLLBACK") | |
| if new_path.exists(): | |
| new_path.unlink() | |
| raise | |
| os.replace(new_path, token_path) | |
| finally: | |
| conn.close() | |
| print(f"Operator token rotated; the old one is revoked. New token saved to {token_path}.") | |
| return 0 | |
| def cmd_backup(args) -> int: | |
| db_path = Path(args.data_dir).expanduser() / DB_NAME | |
| if not db_path.exists(): | |
| print(f"error: no board at {db_path}", file=sys.stderr) | |
| return 1 | |
| conn = dbmod.connect(db_path) | |
| try: | |
| dbmod.backup(conn, args.out) | |
| finally: | |
| conn.close() | |
| print(f"Backup written to {args.out}. Attachments live in {db_path.parent / 'blobs'}; copy them too.") | |
| return 0 | |
| def _is_loopback(host: str) -> bool: | |
| if host == "localhost": | |
| return True | |
| try: | |
| return ipaddress.ip_address(host).is_loopback | |
| except ValueError: | |
| return False | |
| def cmd_serve(args) -> int: | |
| from .server import AlreadyRunning, Board, NotInitialized | |
| data_dir = Path(args.data_dir).expanduser() | |
| _warn_fs(data_dir) | |
| replicator = None | |
| if args.replica: | |
| from .replicate import Replicator, restore, store_from_url | |
| store = store_from_url(args.replica) | |
| if not (data_dir / DB_NAME).exists(): | |
| manifest = restore(store, data_dir) | |
| print(f"Restored {manifest['snapshot']} from replica." if manifest | |
| else "Replica is empty; nothing to restore.", file=sys.stderr) | |
| replicator = Replicator(store, data_dir, interval=args.snapshot_interval) | |
| try: | |
| board = Board(data_dir, replicator=replicator) | |
| except (NotInitialized, AlreadyRunning, dbmod.SchemaError) as exc: | |
| print(f"error: {exc}", file=sys.stderr) | |
| return 1 | |
| if not _is_loopback(args.host): | |
| print(f"warning: listening on {args.host}. Tokens travel in plain HTTP; " | |
| "put HTTPS in front of this outside isolated local development.", file=sys.stderr) | |
| run_board(board, args.host, args.port, args.log_level) | |
| return 0 | |
| def run_board(board, host: str, port: int, log_level: str = "info", app=None) -> None: | |
| """Serve `board` in the foreground until SIGINT/SIGTERM, then close it (final snapshot included).""" | |
| import uvicorn | |
| from .server import create_app | |
| class BoardServer(uvicorn.Server): | |
| def handle_exit(self, sig, frame): | |
| board.begin_shutdown() | |
| super().handle_exit(sig, frame) | |
| config = uvicorn.Config( | |
| app or create_app(board), host=host, port=port, log_level=log_level, | |
| timeout_graceful_shutdown=5, server_header=False, | |
| ) | |
| # uvicorn re-raises the signal it caught after shutting down. With the | |
| # default SIGTERM action that kills the process before board.close() runs, | |
| # so the final snapshot would be lost. Turn it into a normal exit instead. | |
| def _exit(signum, frame): | |
| raise SystemExit(0) | |
| signal.signal(signal.SIGTERM, _exit) | |
| signal.signal(signal.SIGINT, _exit) | |
| try: | |
| BoardServer(config).run() | |
| finally: | |
| board.close() | |
| print("Board stopped cleanly.", file=sys.stderr) | |
| def cmd_restore(args) -> int: | |
| from .replicate import restore, store_from_url | |
| manifest = restore(store_from_url(args.replica), Path(args.data_dir).expanduser()) | |
| if manifest is None: | |
| print("error: replica has no snapshot", file=sys.stderr) | |
| return 1 | |
| print(f"Restored {manifest['snapshot']} ({manifest['size']} bytes) into {args.data_dir}.") | |
| return 0 | |
| def cmd_demo(args) -> int: | |
| from .demo import run_demo | |
| return run_demo(args.data_dir, args.port) | |
| def cmd_deploy_space(args) -> int: | |
| try: | |
| from huggingface_hub import get_token | |
| from huggingface_hub.errors import HfHubHTTPError | |
| except ImportError: | |
| print("error: deploy-space needs the hf extra: pip install 'msgboard[hf]'", file=sys.stderr) | |
| return 1 | |
| from .auth import new_token | |
| from .client import Client | |
| from .space import DeployError, deploy | |
| from .space.hub import _resolve_space_id | |
| try: | |
| space_id = _resolve_space_id(args.space_id) | |
| token_path = Path(args.token_file or Path.home() / ".config/msgboard/spaces" / space_id / OPERATOR_TOKEN_FILE) | |
| token_path = token_path.expanduser() | |
| # A new token replaces the Space's MSGBOARD_OPERATOR_TOKEN secret, which the app adopts on restart. | |
| operator = token_path.read_text().strip() if token_path.exists() else new_token()[1] | |
| result = deploy(space_id, operator_token=operator, bucket_id=args.bucket, private=args.private, | |
| title=args.title) | |
| except (DeployError, HfHubHTTPError, ValueError) as exc: | |
| print(f"error: {exc}", file=sys.stderr) | |
| return 1 | |
| if not token_path.exists(): | |
| token_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) | |
| _write_secret_file(token_path, operator) | |
| print(f"{'Created' if result.created else 'Updated'} the {'private' if result.private else 'public'} " | |
| f"Space {result.space_id} at {result.url}") | |
| print(f"Snapshots go to the private bucket {result.bucket_id}.") | |
| print(f"Operator token in {token_path} (mode 0600). Keep it secret; open the dashboard with it to create invite links.") | |
| if args.wait is None: | |
| print("The Space is building; the board answers in a few minutes.") | |
| return 0 | |
| print(f"Waiting up to {args.wait:g} s for the board to answer ...") | |
| gate = {} | |
| if result.private: # the Space's gate wants an HF token in Authorization; the board token moves aside | |
| gate = {"headers": {"Authorization": f"Bearer {os.environ.get('MSGBOARD_HF_TOKEN') or get_token()}"}, | |
| "token_header": "X-Msgboard-Token"} | |
| with Client(result.url, operator, **gate) as client: | |
| try: | |
| client.wait_ready(args.wait) | |
| except TimeoutError as exc: | |
| print(f"error: {exc}", file=sys.stderr) | |
| return 1 | |
| print(f"Board ready at {result.url}") | |
| return 0 | |
| def main(argv: list[str] | None = None) -> int: | |
| parser = argparse.ArgumentParser(prog="python -m msgboard", description=__doc__) | |
| sub = parser.add_subparsers(dest="command", required=True) | |
| p = sub.add_parser("init", help="create or migrate a board; never overwrites existing credentials") | |
| p.add_argument("--data-dir", required=True) | |
| p.set_defaults(func=cmd_init) | |
| p = sub.add_parser("serve", help="run the HTTP server in the foreground (Ctrl-C to stop)") | |
| p.add_argument("--data-dir", required=True) | |
| p.add_argument("--host", default="127.0.0.1") | |
| p.add_argument("--port", type=int, default=DEFAULT_PORT) | |
| p.add_argument("--log-level", default="info") | |
| p.add_argument("--replica", default=None, | |
| help="snapshot target: hf://buckets/<ns>/<bucket>[/<prefix>] or a directory; " | |
| "an empty data dir is restored from it first") | |
| p.add_argument("--snapshot-interval", type=float, default=10.0, help="seconds between snapshots") | |
| p.set_defaults(func=cmd_serve) | |
| p = sub.add_parser("restore", help="restore an empty data dir from a replica") | |
| p.add_argument("--data-dir", required=True) | |
| p.add_argument("--replica", required=True) | |
| p.set_defaults(func=cmd_restore) | |
| p = sub.add_parser("backup", help="write a consistent SQLite snapshot") | |
| p.add_argument("--data-dir", required=True) | |
| p.add_argument("--out", required=True) | |
| p.set_defaults(func=cmd_backup) | |
| p = sub.add_parser("rotate-operator-token", help="replace the operator token (local file access required)") | |
| p.add_argument("--data-dir", required=True) | |
| p.set_defaults(func=cmd_rotate_operator) | |
| p = sub.add_parser("demo", help="start a throwaway board and run two agents against it") | |
| p.add_argument("--data-dir", default=None, help="defaults to a new temporary directory") | |
| p.add_argument("--port", type=int, default=0, help="0 picks a free port") | |
| p.set_defaults(func=cmd_demo) | |
| p = sub.add_parser("deploy-space", help="create or update a board on a Hugging Face Space (needs the hf extra)") | |
| p.add_argument("space_id", help="namespace/name, or a bare name for the logged-in account") | |
| p.add_argument("--bucket", default=None, | |
| help="snapshot bucket; default: the one the Space uses, else <space_id>-bucket (or -2, -3 ...)") | |
| visibility = p.add_mutually_exclusive_group() | |
| visibility.add_argument("--private", dest="private", action="store_const", const=True, | |
| help="make a new Space private (default: the namespace's default)") | |
| visibility.add_argument("--public", dest="private", action="store_const", const=False, | |
| help="make a new Space public; an existing Space keeps its visibility") | |
| p.add_argument("--title", default=None, help="the Space card's title (default: the Space's name)") | |
| p.add_argument("--token-file", default=None, | |
| help="operator token file, created (0600) if missing; " | |
| "default ~/.config/msgboard/spaces/<namespace>/<name>/operator.token") | |
| p.add_argument("--wait", type=float, nargs="?", const=600.0, default=None, metavar="SECONDS", | |
| help="wait for the board to answer (default 600 s)") | |
| p.set_defaults(func=cmd_deploy_space, private=None) | |
| args = parser.parse_args(argv) | |
| return args.func(args) | |