Merge commit 'f4b3a46109a8da00fdd68d8999160ddc45dcc88a' as 'scripts/agave-container'

scripts/agave-container/entrypoint.py

@@ -0,0 +1,485 @@
+#!/usr/bin/env python3
+"""Agave validator entrypoint — snapshot management, arg construction, liveness probe.
+
+Two subcommands:
+  entrypoint.py serve   (default) — snapshot freshness check + exec agave-validator
+  entrypoint.py probe   — liveness probe (slot lag check, exits 0/1)
+
+Replaces the bash entrypoint.sh / start-rpc.sh / start-validator.sh with a single
+Python module. Test mode still dispatches to start-test.sh.
+
+All configuration comes from environment variables — same vars as the original
+bash scripts. See compose files for defaults.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+import subprocess
+import sys
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+from urllib.request import Request
+
+log: logging.Logger = logging.getLogger("entrypoint")
+
+# Directories
+CONFIG_DIR = "/data/config"
+LEDGER_DIR = "/data/ledger"
+ACCOUNTS_DIR = "/data/accounts"
+SNAPSHOTS_DIR = "/data/snapshots"
+LOG_DIR = "/data/log"
+IDENTITY_FILE = f"{CONFIG_DIR}/validator-identity.json"
+
+# Snapshot filename pattern
+FULL_SNAP_RE: re.Pattern[str] = re.compile(
+    r"^snapshot-(\d+)-[A-Za-z0-9]+\.tar\.(zst|bz2)$"
+)
+
+MAINNET_RPC = "https://api.mainnet-beta.solana.com"
+
+
+# -- Helpers -------------------------------------------------------------------
+
+
+def env(name: str, default: str = "") -> str:
+    """Read env var with default."""
+    return os.environ.get(name, default)
+
+
+def env_required(name: str) -> str:
+    """Read required env var, exit if missing."""
+    val = os.environ.get(name)
+    if not val:
+        log.error("%s is required but not set", name)
+        sys.exit(1)
+    return val
+
+
+def env_bool(name: str, default: bool = False) -> bool:
+    """Read boolean env var (true/false/1/0)."""
+    val = os.environ.get(name, "").lower()
+    if not val:
+        return default
+    return val in ("true", "1", "yes")
+
+
+def rpc_get_slot(url: str, timeout: int = 10) -> int | None:
+    """Get current slot from a Solana RPC endpoint."""
+    payload = json.dumps({
+        "jsonrpc": "2.0", "id": 1,
+        "method": "getSlot", "params": [],
+    }).encode()
+    req = Request(url, data=payload,
+                  headers={"Content-Type": "application/json"})
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            data = json.loads(resp.read())
+            result = data.get("result")
+            if isinstance(result, int):
+                return result
+    except (urllib.error.URLError, json.JSONDecodeError, OSError, TimeoutError):
+        pass
+    return None
+
+
+# -- Snapshot management -------------------------------------------------------
+
+
+def get_local_snapshot_slot(snapshots_dir: str) -> int | None:
+    """Find the highest slot among local snapshot files."""
+    best_slot: int | None = None
+    snap_path = Path(snapshots_dir)
+    if not snap_path.is_dir():
+        return None
+    for entry in snap_path.iterdir():
+        m = FULL_SNAP_RE.match(entry.name)
+        if m:
+            slot = int(m.group(1))
+            if best_slot is None or slot > best_slot:
+                best_slot = slot
+    return best_slot
+
+
+def clean_snapshots(snapshots_dir: str) -> None:
+    """Remove all snapshot files from the directory."""
+    snap_path = Path(snapshots_dir)
+    if not snap_path.is_dir():
+        return
+    for entry in snap_path.iterdir():
+        if entry.name.startswith(("snapshot-", "incremental-snapshot-")):
+            log.info("Removing old snapshot: %s", entry.name)
+            entry.unlink(missing_ok=True)
+
+
+def maybe_download_snapshot(snapshots_dir: str) -> None:
+    """Check snapshot freshness and download if needed.
+
+    Controlled by env vars:
+      SNAPSHOT_AUTO_DOWNLOAD (default: true) — enable/disable
+      SNAPSHOT_MAX_AGE_SLOTS (default: 20000) — staleness threshold
+    """
+    if not env_bool("SNAPSHOT_AUTO_DOWNLOAD", default=True):
+        log.info("Snapshot auto-download disabled")
+        return
+
+    max_age = int(env("SNAPSHOT_MAX_AGE_SLOTS", "20000"))
+
+    # Get mainnet current slot
+    mainnet_slot = rpc_get_slot(MAINNET_RPC)
+    if mainnet_slot is None:
+        log.warning("Cannot reach mainnet RPC — skipping snapshot check")
+        return
+
+    # Check local snapshot
+    local_slot = get_local_snapshot_slot(snapshots_dir)
+    if local_slot is not None:
+        age = mainnet_slot - local_slot
+        log.info("Local snapshot at slot %d (mainnet: %d, age: %d slots)",
+                 local_slot, mainnet_slot, age)
+        if age <= max_age:
+            log.info("Snapshot is fresh enough (age %d <= %d), skipping download", age, max_age)
+            return
+        log.info("Snapshot is stale (age %d > %d), downloading fresh", age, max_age)
+    else:
+        log.info("No local snapshot found, downloading")
+
+    # Clean old snapshots before downloading
+    clean_snapshots(snapshots_dir)
+
+    # Import and call snapshot download
+    # snapshot_download.py is installed alongside this file in /usr/local/bin/
+    script_dir = Path(__file__).resolve().parent
+    sys.path.insert(0, str(script_dir))
+    from snapshot_download import download_best_snapshot
+
+    ok = download_best_snapshot(snapshots_dir)
+    if not ok:
+        log.error("Snapshot download failed — starting without fresh snapshot")
+
+
+# -- Directory and identity setup ----------------------------------------------
+
+
+def ensure_dirs(*dirs: str) -> None:
+    """Create directories and fix ownership."""
+    uid = os.getuid()
+    gid = os.getgid()
+    for d in dirs:
+        os.makedirs(d, exist_ok=True)
+        try:
+            subprocess.run(
+                ["sudo", "chown", "-R", f"{uid}:{gid}", d],
+                check=False, capture_output=True,
+            )
+        except FileNotFoundError:
+            pass  # sudo not available — dirs already owned correctly
+
+
+def ensure_identity_rpc() -> None:
+    """Generate ephemeral identity keypair for RPC mode if not mounted."""
+    if os.path.isfile(IDENTITY_FILE):
+        return
+    log.info("Generating RPC node identity keypair...")
+    subprocess.run(
+        ["solana-keygen", "new", "--no-passphrase", "--silent",
+         "--force", "--outfile", IDENTITY_FILE],
+        check=True,
+    )
+
+
+def print_identity() -> None:
+    """Print the node identity pubkey."""
+    result = subprocess.run(
+        ["solana-keygen", "pubkey", IDENTITY_FILE],
+        capture_output=True, text=True, check=False,
+    )
+    if result.returncode == 0:
+        log.info("Node identity: %s", result.stdout.strip())
+
+
+# -- Arg construction ----------------------------------------------------------
+
+
+def build_common_args() -> list[str]:
+    """Build agave-validator args common to both RPC and validator modes."""
+    args: list[str] = [
+        "--identity", IDENTITY_FILE,
+        "--entrypoint", env_required("VALIDATOR_ENTRYPOINT"),
+        "--known-validator", env_required("KNOWN_VALIDATOR"),
+        "--ledger", LEDGER_DIR,
+        "--accounts", ACCOUNTS_DIR,
+        "--snapshots", SNAPSHOTS_DIR,
+        "--rpc-port", env("RPC_PORT", "8899"),
+        "--rpc-bind-address", env("RPC_BIND_ADDRESS", "127.0.0.1"),
+        "--gossip-port", env("GOSSIP_PORT", "8001"),
+        "--dynamic-port-range", env("DYNAMIC_PORT_RANGE", "9000-10000"),
+        "--no-os-network-limits-test",
+        "--wal-recovery-mode", "skip_any_corrupted_record",
+        "--limit-ledger-size", env("LIMIT_LEDGER_SIZE", "50000000"),
+    ]
+
+    # Snapshot generation
+    if env("NO_SNAPSHOTS") == "true":
+        args.append("--no-snapshots")
+    else:
+        args += [
+            "--full-snapshot-interval-slots", env("SNAPSHOT_INTERVAL_SLOTS", "100000"),
+            "--maximum-full-snapshots-to-retain", env("MAXIMUM_SNAPSHOTS_TO_RETAIN", "5"),
+        ]
+        if env("NO_INCREMENTAL_SNAPSHOTS") != "true":
+            args += ["--maximum-incremental-snapshots-to-retain", "2"]
+
+    # Account indexes
+    account_indexes = env("ACCOUNT_INDEXES")
+    if account_indexes:
+        for idx in account_indexes.split(","):
+            idx = idx.strip()
+            if idx:
+                args += ["--account-index", idx]
+
+    # Additional entrypoints
+    for ep in env("EXTRA_ENTRYPOINTS").split():
+        if ep:
+            args += ["--entrypoint", ep]
+
+    # Additional known validators
+    for kv in env("EXTRA_KNOWN_VALIDATORS").split():
+        if kv:
+            args += ["--known-validator", kv]
+
+    # Cluster verification
+    genesis_hash = env("EXPECTED_GENESIS_HASH")
+    if genesis_hash:
+        args += ["--expected-genesis-hash", genesis_hash]
+    shred_version = env("EXPECTED_SHRED_VERSION")
+    if shred_version:
+        args += ["--expected-shred-version", shred_version]
+
+    # Metrics — just needs to be in the environment, agave reads it directly
+    # (env var is already set, nothing to pass as arg)
+
+    # Gossip host / TVU address
+    gossip_host = env("GOSSIP_HOST")
+    if gossip_host:
+        args += ["--gossip-host", gossip_host]
+    elif env("PUBLIC_TVU_ADDRESS"):
+        args += ["--public-tvu-address", env("PUBLIC_TVU_ADDRESS")]
+
+    # Jito flags
+    if env("JITO_ENABLE") == "true":
+        log.info("Jito MEV enabled")
+        jito_flags: list[tuple[str, str]] = [
+            ("JITO_TIP_PAYMENT_PROGRAM", "--tip-payment-program-pubkey"),
+            ("JITO_DISTRIBUTION_PROGRAM", "--tip-distribution-program-pubkey"),
+            ("JITO_MERKLE_ROOT_AUTHORITY", "--merkle-root-upload-authority"),
+            ("JITO_COMMISSION_BPS", "--commission-bps"),
+            ("JITO_BLOCK_ENGINE_URL", "--block-engine-url"),
+            ("JITO_SHRED_RECEIVER_ADDR", "--shred-receiver-address"),
+        ]
+        for env_name, flag in jito_flags:
+            val = env(env_name)
+            if val:
+                args += [flag, val]
+
+    return args
+
+
+def build_rpc_args() -> list[str]:
+    """Build agave-validator args for RPC (non-voting) mode."""
+    args = build_common_args()
+    args += [
+        "--no-voting",
+        "--log", f"{LOG_DIR}/validator.log",
+        "--full-rpc-api",
+        "--enable-rpc-transaction-history",
+        "--rpc-pubsub-enable-block-subscription",
+        "--enable-extended-tx-metadata-storage",
+        "--no-wait-for-vote-to-start-leader",
+        "--no-snapshot-fetch",
+    ]
+
+    # Public vs private RPC
+    public_rpc = env("PUBLIC_RPC_ADDRESS")
+    if public_rpc:
+        args += ["--public-rpc-address", public_rpc]
+    else:
+        args += ["--private-rpc", "--allow-private-addr", "--only-known-rpc"]
+
+    # Jito relayer URL (RPC mode doesn't use it, but validator mode does —
+    # handled in build_validator_args)
+
+    return args
+
+
+def build_validator_args() -> list[str]:
+    """Build agave-validator args for voting validator mode."""
+    vote_keypair = env("VOTE_ACCOUNT_KEYPAIR",
+                       "/data/config/vote-account-keypair.json")
+
+    # Identity must be mounted for validator mode
+    if not os.path.isfile(IDENTITY_FILE):
+        log.error("Validator identity keypair not found at %s", IDENTITY_FILE)
+        log.error("Mount your validator keypair to %s", IDENTITY_FILE)
+        sys.exit(1)
+
+    # Vote account keypair must exist
+    if not os.path.isfile(vote_keypair):
+        log.error("Vote account keypair not found at %s", vote_keypair)
+        log.error("Mount your vote account keypair or set VOTE_ACCOUNT_KEYPAIR")
+        sys.exit(1)
+
+    # Print vote account pubkey
+    result = subprocess.run(
+        ["solana-keygen", "pubkey", vote_keypair],
+        capture_output=True, text=True, check=False,
+    )
+    if result.returncode == 0:
+        log.info("Vote account: %s", result.stdout.strip())
+
+    args = build_common_args()
+    args += [
+        "--vote-account", vote_keypair,
+        "--log", "-",
+    ]
+
+    # Jito relayer URL (validator-only)
+    relayer_url = env("JITO_RELAYER_URL")
+    if env("JITO_ENABLE") == "true" and relayer_url:
+        args += ["--relayer-url", relayer_url]
+
+    return args
+
+
+def append_extra_args(args: list[str]) -> list[str]:
+    """Append EXTRA_ARGS passthrough flags."""
+    extra = env("EXTRA_ARGS")
+    if extra:
+        args += extra.split()
+    return args
+
+
+# -- Serve subcommand ---------------------------------------------------------
+
+
+def cmd_serve() -> None:
+    """Main serve flow: snapshot check, setup, exec agave-validator."""
+    mode = env("AGAVE_MODE", "test")
+    log.info("AGAVE_MODE=%s", mode)
+
+    # Test mode dispatches to start-test.sh
+    if mode == "test":
+        os.execvp("start-test.sh", ["start-test.sh"])
+
+    if mode not in ("rpc", "validator"):
+        log.error("Unknown AGAVE_MODE: %s (valid: test, rpc, validator)", mode)
+        sys.exit(1)
+
+    # Ensure directories
+    dirs = [CONFIG_DIR, LEDGER_DIR, ACCOUNTS_DIR, SNAPSHOTS_DIR]
+    if mode == "rpc":
+        dirs.append(LOG_DIR)
+    ensure_dirs(*dirs)
+
+    # Snapshot freshness check and auto-download
+    maybe_download_snapshot(SNAPSHOTS_DIR)
+
+    # Identity setup
+    if mode == "rpc":
+        ensure_identity_rpc()
+    print_identity()
+
+    # Build args
+    if mode == "rpc":
+        args = build_rpc_args()
+    else:
+        args = build_validator_args()
+
+    args = append_extra_args(args)
+
+    # Write startup timestamp for probe grace period
+    Path("/tmp/entrypoint-start").write_text(str(time.time()))
+
+    log.info("Starting agave-validator with %d arguments", len(args))
+    os.execvp("agave-validator", ["agave-validator"] + args)
+
+
+# -- Probe subcommand ---------------------------------------------------------
+
+
+def cmd_probe() -> None:
+    """Liveness probe: check local RPC slot vs mainnet.
+
+    Exit 0 = healthy, exit 1 = unhealthy.
+
+    Grace period: PROBE_GRACE_SECONDS (default 600) — probe always passes
+    during grace period to allow for snapshot unpacking and initial replay.
+    """
+    grace_seconds = int(env("PROBE_GRACE_SECONDS", "600"))
+    max_lag = int(env("PROBE_MAX_SLOT_LAG", "20000"))
+
+    # Check grace period
+    start_file = Path("/tmp/entrypoint-start")
+    if start_file.exists():
+        try:
+            start_time = float(start_file.read_text().strip())
+            elapsed = time.time() - start_time
+            if elapsed < grace_seconds:
+                # Within grace period — always healthy
+                sys.exit(0)
+        except (ValueError, OSError):
+            pass
+    else:
+        # No start file — serve hasn't started yet, within grace
+        sys.exit(0)
+
+    # Query local RPC
+    rpc_port = env("RPC_PORT", "8899")
+    local_url = f"http://127.0.0.1:{rpc_port}"
+    local_slot = rpc_get_slot(local_url, timeout=5)
+    if local_slot is None:
+        # Local RPC unreachable after grace period — unhealthy
+        sys.exit(1)
+
+    # Query mainnet
+    mainnet_slot = rpc_get_slot(MAINNET_RPC, timeout=10)
+    if mainnet_slot is None:
+        # Can't reach mainnet to compare — assume healthy (don't penalize
+        # the validator for mainnet RPC being down)
+        sys.exit(0)
+
+    lag = mainnet_slot - local_slot
+    if lag > max_lag:
+        sys.exit(1)
+
+    sys.exit(0)
+
+
+# -- Main ----------------------------------------------------------------------
+
+
+def main() -> None:
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
+        datefmt="%H:%M:%S",
+    )
+
+    subcmd = sys.argv[1] if len(sys.argv) > 1 else "serve"
+
+    if subcmd == "serve":
+        cmd_serve()
+    elif subcmd == "probe":
+        cmd_probe()
+    else:
+        log.error("Unknown subcommand: %s (valid: serve, probe)", subcmd)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()