path: root/scripts/test_scripts/hub/seed_s3_snapshot.py

#!/usr/bin/env python3
"""Stage A of the perf-seed workflow.

Fetches real module data from a configured S3 bucket via the hub hydration
pipeline and preserves the decompressed server-state trees to disk so a
later stage can replay them into a local MinIO backend.

Flow:
  1. Start a hub against real S3 (read-only SSO credentials), dehydration off,
     hibernated-watchdog timeout set huge so modules are not auto-deprovisioned
     mid-flow.
  2. Provision N modules (first N moduleIds from the bucket listing, filtered
     to UUID-shaped folders).
  3. Wait until every module reaches 'provisioned'.
  4. Hibernate every module (child processes exit, state dir intact and no
     remaining writer while we copy).
  5. Wait until every module reaches 'hibernated'.
  6. Copy <hub-data-dir>/servers/<moduleid>/ -> <snapshot-dir>/<moduleid>/
     with the hub still running - hibernate took the writers offline and the
     watchdog is muted.
  7. Shut the hub down via 'zen down --pid <hub pid>'.

Required environment variables (or pass via CLI flags):
  ZEN_PERF_S3_URI       e.g. s3://your-bucket/optional-prefix/
  ZEN_PERF_AWS_PROFILE  AWS SSO profile name configured with read access
  ZEN_PERF_AWS_REGION   defaults to us-east-1

Requirements:
  pip install boto3
  aws CLI v2 with the SSO profile configured
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Optional

_EXE_SUFFIX = ".exe" if sys.platform == "win32" else ""


def _rmtree_robust(path) -> None:
    """shutil.rmtree with a Windows-friendly retry for read-only files."""
    import os as _os
    import stat as _stat
    def _onerror(func, p, exc_info):
        try:
            _os.chmod(p, _stat.S_IWRITE)
            func(p)
        except Exception:
            pass
    # onexc was introduced in py3.12; fall back to onerror on older versions.
    if sys.version_info >= (3, 12):
        shutil.rmtree(path, onexc=lambda func, p, exc: _onerror(func, p, (type(exc), exc, exc.__traceback__)))
    else:
        shutil.rmtree(path, onerror=_onerror)


_DEFAULT_S3_URI = os.environ.get("ZEN_PERF_S3_URI", "")
_DEFAULT_AWS_PROFILE = os.environ.get("ZEN_PERF_AWS_PROFILE", "")
_DEFAULT_AWS_REGION = os.environ.get("ZEN_PERF_AWS_REGION", "us-east-1")

# Matches UUID-shaped module IDs (UUID-ish with mixed 4-char groups). Filters
# out stray top-level keys (readmes, test scratches) that aren't real modules.
_MODULEID_RE = re.compile(r"^[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}$")


# ---------------------------------------------------------------------------
# zenserver binary discovery - prefer release, fall back to debug.
# ---------------------------------------------------------------------------

def _find_zenserver(override: Optional[str]) -> Path:
    if override:
        p = Path(override) / f"zenserver{_EXE_SUFFIX}"
        if not p.exists():
            sys.exit(f"zenserver not found at {p}")
        return p

    script_dir = Path(__file__).resolve().parent
    repo_root = script_dir.parents[2]
    for mode in ("release", "debug"):
        for plat in (("windows", "x64"), ("linux", "x86_64"), ("macosx", "x86_64")):
            p = repo_root / "build" / plat[0] / plat[1] / mode / f"zenserver{_EXE_SUFFIX}"
            if p.exists():
                return p
    sys.exit("zenserver executable not found under build/. Build it or pass --zenserver-dir.")


def _find_zen(zenserver_exe: Path) -> Path:
    p = zenserver_exe.parent / f"zen{_EXE_SUFFIX}"
    if not p.exists():
        sys.exit(f"zen CLI not found at {p} (used for graceful hub shutdown)")
    return p


# ---------------------------------------------------------------------------
# AWS SSO credential resolution
# ---------------------------------------------------------------------------

def _require_boto3():
    try:
        import boto3  # type: ignore[import-not-found]
        import botocore.exceptions  # type: ignore[import-not-found]
    except ImportError:
        sys.exit("[aws] boto3 is required. Install it with: pip install boto3")
    return boto3, botocore.exceptions


def _sso_login(profile: str) -> None:
    print(f"[aws] running 'aws sso login --profile {profile}'...")
    rc = subprocess.call(["aws", "sso", "login", "--profile", profile])
    if rc != 0:
        sys.exit(f"[aws] 'aws sso login' failed with rc={rc}")


def _get_session(profile: str):
    boto3, exc_mod = _require_boto3()

    def _load_frozen():
        session = boto3.Session(profile_name=profile)
        creds = session.get_credentials()
        if creds is None:
            return session, None
        return session, creds.get_frozen_credentials()

    try:
        session, frozen = _load_frozen()
        if frozen is not None and frozen.access_key and frozen.secret_key:
            return session, frozen
    except exc_mod.ProfileNotFound:
        sys.exit(f"[aws] profile '{profile}' not found in ~/.aws/config")
    except Exception as e:
        print(f"[aws] initial credential load failed: {e}")

    _sso_login(profile)

    session, frozen = _load_frozen()
    if frozen is None or not frozen.access_key:
        sys.exit("[aws] could not resolve credentials after sso login")
    return session, frozen


def _parse_s3_uri(uri: str) -> tuple[str, str]:
    if not uri.startswith("s3://"):
        sys.exit(f"[aws] invalid S3 URI: {uri}")
    rest = uri[len("s3://"):]
    if "/" in rest:
        bucket, prefix = rest.split("/", 1)
    else:
        bucket, prefix = rest, ""
    return bucket, prefix


_EMPTY_MIN_OBJECTS = 3
_EMPTY_MIN_BYTES = 16 * 1024

# Size buckets for stratified selection. Pyramid distribution mirrors typical
# workloads: many small modules, fewer medium, fewest large.
_SMALL_MAX_BYTES = 1 * 1024 * 1024            # <1 MiB
_LARGE_MIN_BYTES = 500 * 1024 * 1024          # >500 MiB
_BUCKET_RATIO = (("small", 500), ("medium", 350), ("large", 150))


def _list_module_ids(session, bucket: str, prefix: str, region: str, limit: int) -> list[str]:
    """List UUID-shaped module folders under the bucket and return the `limit`
    most recently active non-empty ones, ranked by the LastModified of each
    module's `incremental-state.cbo` (newest first - these were last dehydrated
    most recently, which is the closest proxy for "most recently accessed").

    Modules with < _EMPTY_MIN_OBJECTS objects or < _EMPTY_MIN_BYTES total bytes
    are treated as empty and dropped.
    """
    s3 = session.client("s3", region_name=region)
    prefix_norm = prefix if (not prefix or prefix.endswith("/")) else prefix + "/"

    # 1. Enumerate every UUID-shaped folder.
    paginator = s3.get_paginator("list_objects_v2")
    candidates: list[str] = []
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix_norm, Delimiter="/"):
        for cp in page.get("CommonPrefixes", []) or []:
            folder = cp.get("Prefix", "")[len(prefix_norm):].rstrip("/")
            if folder and _MODULEID_RE.match(folder):
                candidates.append(folder)
    print(f"[s3] {len(candidates)} module folders match UUID shape; "
          f"sizing + ranking by state.cbo LastModified (skip empty <{_EMPTY_MIN_OBJECTS} obj or <{_EMPTY_MIN_BYTES} B)...")

    # 2. Per-module ListObjectsV2: total object count + total bytes + state.cbo
    #    LastModified, all in one pass. Missing state file => epoch sentinel.
    from datetime import datetime, timezone
    epoch = datetime(1970, 1, 1, tzinfo=timezone.utc)

    def _probe(mid: str) -> tuple[int, int, datetime]:
        mid_prefix = f"{prefix_norm}{mid}/"
        state_key = f"{mid_prefix}incremental-state.cbo"
        count = 0
        total_bytes = 0
        state_mtime = epoch
        try:
            for page in paginator.paginate(Bucket=bucket, Prefix=mid_prefix):
                for obj in page.get("Contents", []) or []:
                    count += 1
                    total_bytes += int(obj.get("Size", 0))
                    if obj.get("Key") == state_key:
                        state_mtime = obj.get("LastModified", epoch) or epoch
        except Exception:
            return 0, 0, epoch
        return count, total_bytes, state_mtime

    with ThreadPoolExecutor(max_workers=50) as pool:
        probes = list(pool.map(_probe, candidates))

    # 3. Stratified selection. Bucket non-empty modules by total_bytes into
    #    small/medium/large, sort each bucket by state.cbo mtime desc, and pull
    #    according to _BUCKET_RATIO scaled to `limit`. If a bucket is short,
    #    spill the deficit onto the relaxed pool first, then onto neighbouring
    #    buckets so the caller still gets `limit` modules.
    small: list[tuple[str, datetime, int]] = []
    medium: list[tuple[str, datetime, int]] = []
    large: list[tuple[str, datetime, int]] = []
    relaxed: list[tuple[str, datetime, int, int]] = []
    empties = 0
    no_state = 0
    for mid, (count, total_bytes, mtime) in zip(candidates, probes):
        is_empty = count < _EMPTY_MIN_OBJECTS or total_bytes < _EMPTY_MIN_BYTES
        if is_empty:
            empties += 1
        if mtime == epoch:
            no_state += 1
        if not is_empty and mtime != epoch:
            if total_bytes < _SMALL_MAX_BYTES:
                small.append((mid, mtime, total_bytes))
            elif total_bytes >= _LARGE_MIN_BYTES:
                large.append((mid, mtime, total_bytes))
            else:
                medium.append((mid, mtime, total_bytes))
        else:
            relaxed.append((mid, mtime, count, total_bytes))

    for bucket in (small, medium, large):
        bucket.sort(key=lambda x: x[1], reverse=True)
    print(f"[s3] candidates={len(candidates)} empty={empties} no-state={no_state} "
          f"small={len(small)} medium={len(medium)} large={len(large)}")

    total_ratio = sum(w for _, w in _BUCKET_RATIO)
    targets = {name: max(0, (limit * w) // total_ratio) for name, w in _BUCKET_RATIO}
    # Distribute rounding leftovers to first bucket(s) deterministically.
    leftover = limit - sum(targets.values())
    for name, _ in _BUCKET_RATIO:
        if leftover <= 0:
            break
        targets[name] += 1
        leftover -= 1

    pools = {"small": small, "medium": medium, "large": large}
    # If a bucket is short of its target, the deficit redirects to the nearest
    # neighbour bucket(s) before falling back to anything further away. Small
    # is rare in real workloads, so its deficit pads from medium - not large.
    fallback_chain = {
        "small": ["medium", "large"],
        "medium": ["small", "large"],
        "large": ["medium", "small"],
    }
    cursors = {"small": 0, "medium": 0, "large": 0}
    selected: list[str] = []
    seen: set[str] = set()
    fills = {"small": 0, "medium": 0, "large": 0}
    redirected = {"small": 0, "medium": 0, "large": 0}

    def _take_from(name: str, n: int) -> int:
        """Take up to n entries from pools[name] starting at cursor. Returns
        the number actually taken."""
        if n <= 0:
            return 0
        pool = pools[name]
        start = cursors[name]
        end = start + n
        taken = 0
        for mid, _, _ in pool[start:end]:
            if mid in seen:
                continue
            seen.add(mid)
            selected.append(mid)
            taken += 1
        cursors[name] = end
        return taken

    for name, _ in _BUCKET_RATIO:
        want = targets[name]
        got = _take_from(name, want)
        fills[name] = got
        deficit = want - got
        if deficit <= 0:
            continue
        for fb in fallback_chain[name]:
            if deficit <= 0:
                break
            avail = max(0, len(pools[fb]) - cursors[fb])
            grab = min(deficit, avail)
            if grab <= 0:
                continue
            got_fb = _take_from(fb, grab)
            redirected[name] += got_fb
            deficit -= got_fb
    print(f"[s3] bucket fills: small={fills['small']}+{redirected['small']} "
          f"medium={fills['medium']}+{redirected['medium']} "
          f"large={fills['large']}+{redirected['large']} "
          f"(numbers after '+' = redirected from neighbour buckets)")

    deficit = limit - len(selected)
    if deficit > 0:
        # Last-resort spill: relaxed pool (empties + no-state).
        relaxed.sort(key=lambda x: (x[1], x[3], x[2]), reverse=True)
        added = 0
        for mid, _, _, _ in relaxed:
            if added >= deficit:
                break
            if mid in seen:
                continue
            seen.add(mid)
            selected.append(mid)
            added += 1
        print(f"[s3] extended with {added} relaxed entries to fill deficit")

    return selected[:limit]


# ---------------------------------------------------------------------------
# Hub lifecycle
# ---------------------------------------------------------------------------

def _start_hub(
    zenserver_exe: Path,
    data_dir: Path,
    port: int,
    log_file: Path,
    instance_limit: int,
    extra_args: list[str],
    extra_env: dict[str, str],
) -> tuple[subprocess.Popen, object]:
    data_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        str(zenserver_exe),
        "hub",
        "--enable-execution-history=false",
        f"--data-dir={data_dir}",
        f"--port={port}",
        "--hub-instance-corelimit=4",
        "--hub-provision-disk-limit-percent=99",
        "--hub-provision-memory-limit-percent=80",
        f"--hub-instance-limit={instance_limit}",
        # Provision pool + async S3 in-flight cap use server defaults; on a
        # 128-core host these resolve to 16 / 512 which is what the seed run
        # needs (clamp(cpu/8, 4, 16) and clamp(cpu*4, 128, 512)).
        # With 1000 modules the seeding flow runs for 20+ minutes. Extend BOTH
        # watchdog inactivity timers so early-provisioned modules do not get
        # auto-deprovisioned while we're still hydrating the tail (hibernated
        # default 1800s, provisioned default 600s - the latter is the one that
        # bit us on the 1000-module run and dropped 335 modules).
        "--hub-watchdog-provisioned-inactivity-timeout-seconds=86400",
        "--hub-watchdog-hibernated-inactivity-timeout-seconds=86400",
    ] + extra_args

    env = os.environ.copy()
    env.update(extra_env)

    popen_kwargs: dict = {}
    if sys.platform == "win32":
        popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
    log_handle = log_file.open("wb")
    try:
        proc = subprocess.Popen(
            cmd, env=env, stdout=log_handle, stderr=subprocess.STDOUT,
            **popen_kwargs,
        )
    except Exception:
        log_handle.close()
        raise
    print(f"[hub] started (pid {proc.pid}), log: {log_file}")
    return proc, log_handle


def _wait_for_hub(proc: subprocess.Popen, port: int, timeout_s: float = 100.0) -> None:
    deadline = time.monotonic() + timeout_s
    req = urllib.request.Request(f"http://localhost:{port}/hub/status",
                                 headers={"Accept": "application/json"})
    while time.monotonic() < deadline:
        if proc.poll() is not None:
            sys.exit(f"[hub] process exited unexpectedly (rc={proc.returncode}) - "
                     f"is another zenserver already running on port {port}?")
        try:
            with urllib.request.urlopen(req, timeout=2):
                print("[hub] ready")
                return
        except Exception:
            time.sleep(0.2)
    sys.exit(f"[hub] timed out waiting for readiness after {timeout_s}s")


def _zen_down_hub(zen_exe: Path, hub_proc: subprocess.Popen, timeout_s: float = 300.0) -> None:
    """Shut the hub down via 'zen down --pid <pid>'. zen down performs a
    proper server-initiated shutdown (signals the shutdown event, waits for
    the server to drain) rather than a blunt kill."""
    if hub_proc.poll() is not None:
        return
    pid = hub_proc.pid
    print(f"[hub] zen down --pid {pid}")
    rc = subprocess.call([str(zen_exe), "down", "--pid", str(pid), "--force"])
    if rc != 0:
        print(f"[hub] zen down returned rc={rc}; waiting for exit anyway")
    try:
        hub_proc.wait(timeout=timeout_s)
    except subprocess.TimeoutExpired:
        print(f"[hub] did not exit after {timeout_s}s, killing")
        hub_proc.kill()
        hub_proc.wait()


# ---------------------------------------------------------------------------
# Hub HTTP API
# ---------------------------------------------------------------------------

def _hub_post(port: int, path: str, timeout_s: float = 60.0) -> tuple[int, dict]:
    url = f"http://localhost:{port}{path}"
    req = urllib.request.Request(url, data=b"{}", method="POST",
                                 headers={"Content-Type": "application/json",
                                          "Accept": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            try:
                body = json.loads(resp.read())
            except Exception:
                body = {}
            return resp.status, body
    except urllib.error.HTTPError as e:
        try:
            body = json.loads(e.read())
        except Exception:
            body = {}
        return e.code, body
    except Exception as e:
        return 0, {"error": str(e)}


_HYDRATION_FAIL_RE = re.compile(r"Hydration of module '([0-9a-f-]+)' failed: (.+?)(?:\.\.|$)")


def _scan_hydration_failures(log_path: Path, offset: int, already_warned: set[str], label: str) -> int:
    """Read new hub.log content from `offset` to EOF; print a [warn] line for
    each newly observed `Hydration of module 'X' failed` warning. Returns the
    new offset (resume point for next call)."""
    try:
        size = log_path.stat().st_size
    except OSError:
        return offset
    if size <= offset:
        return offset
    try:
        with open(log_path, "rb") as f:
            f.seek(offset)
            chunk = f.read(size - offset).decode("utf-8", errors="replace")
    except OSError:
        return offset
    for m in _HYDRATION_FAIL_RE.finditer(chunk):
        mid, reason = m.group(1), m.group(2).strip()
        if mid in already_warned:
            continue
        already_warned.add(mid)
        # Trim long S3 error reasons so the line stays scannable.
        if len(reason) > 200:
            reason = reason[:200] + "..."
        print(f"\n[{label}] HYDRATION FAILED {mid}: {reason}", flush=True)
    return size


def _hub_module_states(port: int, timeout_s: float = 10.0) -> Optional[dict[str, str]]:
    url = f"http://localhost:{port}/hub/status"
    req = urllib.request.Request(url, headers={"Accept": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            data = json.loads(resp.read())
    except Exception:
        return None
    return {m["moduleId"]: m.get("state", "") for m in (data.get("modules") or []) if m.get("moduleId")}


def _fan_out_post(
    pool: ThreadPoolExecutor,
    port: int,
    module_ids: list[str],
    verb: str,
) -> tuple[list[str], list[tuple[str, int, dict]]]:
    """POST /hub/modules/<id>/<verb> concurrently. Returns (accepted, failures)."""
    futures = {
        mid: pool.submit(_hub_post, port, f"/hub/modules/{mid}/{verb}")
        for mid in module_ids
    }
    accepted: list[str] = []
    failures: list[tuple[str, int, dict]] = []
    for mid, fut in futures.items():
        status, body = fut.result()
        if status in (200, 202):
            accepted.append(mid)
        else:
            failures.append((mid, status, body))
    return accepted, failures


# Any of these means the module is done transitioning and will not reach the
# target state on its own (without a retry we control).
_FAILED_STATES = {"crashed", "unprovisioned"}


def _wait_for_state(
    port: int,
    module_ids: list[str],
    target_state: str,
    timeout_s: float,
    label: str,
    hub_log: Optional[Path] = None,
    hydration_warned: Optional[set[str]] = None,
) -> tuple[list[str], list[str], dict[str, str]]:
    """Poll hub status until every module hits target_state, fails, or times out.

    Returns (stuck, failed, last_states). 'stuck' = still mid-transition when
    we timed out. 'failed' = hit an _FAILED_STATES value such as 'crashed'.

    When hub_log + hydration_warned are supplied, the wait loop also tails the
    hub log and surfaces `Hydration of module ... failed` warnings as they
    appear. Hydration failures do NOT push modules into _FAILED_STATES (the hub
    cleans the state and still marks the module 'provisioned'), so this is the
    only way to see them in script output.
    """
    deadline = time.monotonic() + timeout_s
    remaining = set(module_ids)
    failed: list[str] = []
    last_states: dict[str, str] = {mid: "" for mid in module_ids}
    log_offset = hub_log.stat().st_size if (hub_log and hub_log.exists()) else 0

    while remaining and time.monotonic() < deadline:
        states = _hub_module_states(port)
        if states is not None:
            for mid in list(remaining):
                s = states.get(mid, "")
                last_states[mid] = s
                if s == target_state:
                    remaining.discard(mid)
                elif s in _FAILED_STATES and target_state not in _FAILED_STATES:
                    remaining.discard(mid)
                    failed.append(mid)
        done = len(module_ids) - len(remaining)
        warned_count = len(hydration_warned) if hydration_warned is not None else 0
        suffix = f", hydration-failed {warned_count}" if hydration_warned is not None else ""
        print(f"[{label}] {done}/{len(module_ids)} '{target_state}' ({len(failed)} failed{suffix})...", end="\r")
        if hub_log is not None and hydration_warned is not None:
            log_offset = _scan_hydration_failures(hub_log, log_offset, hydration_warned, label)
        time.sleep(2.0)

    print()
    if hub_log is not None and hydration_warned is not None:
        _scan_hydration_failures(hub_log, log_offset, hydration_warned, label)
    return list(remaining), failed, last_states


# ---------------------------------------------------------------------------
# Snapshot copy (run AFTER hub shutdown so there's no concurrent writer)
# ---------------------------------------------------------------------------

def _same_volume(a: Path, b: Path) -> bool:
    """True when a and b live on the same filesystem volume (so an os.replace
    rename is O(1) instead of an EXDEV-driven copy+delete fallback)."""
    try:
        return os.stat(a).st_dev == os.stat(b).st_dev
    except OSError:
        return False


def _copy_snapshot(src_root: Path, dst_root: Path, module_ids: list[str]) -> tuple[int, int, int]:
    """Move (preferred) or copy src_root/<mid>/ to dst_root/<mid>/ for each module.

    Hub data dir is throwaway after Stage A (the hub is shut down right after
    this step), so per-module trees can be moved out of it. When src_root and
    dst_root share a volume the move is an O(1) directory rename; when they
    don't, shutil.move falls back to a byte copy plus rmtree, which matches
    the old shutil.copytree cost.

    Returns (modules_moved, files_moved, bytes_moved).
    Replaces only the specific per-module subdirs; never touches siblings.
    """
    dst_root.mkdir(parents=True, exist_ok=True)
    same_vol = _same_volume(src_root, dst_root)
    if same_vol:
        print(f"[snapshot] src and dst share a volume; moving per-module trees (O(1) rename per module)")
    else:
        print(f"[snapshot] src and dst are on different volumes; falling back to byte copy")

    modules_moved = 0
    files_moved = 0
    bytes_moved = 0

    for i, mid in enumerate(module_ids, 1):
        src = src_root / mid
        dst = dst_root / mid
        if not src.is_dir():
            print(f"[snapshot] WARNING: source missing for {mid}: {src}")
            continue
        if dst.exists():
            _rmtree_robust(dst)
        if same_vol:
            os.replace(src, dst)
        else:
            shutil.move(str(src), str(dst))
        modules_moved += 1
        for root, _dirs, files in os.walk(dst):
            for f in files:
                p = Path(root) / f
                try:
                    bytes_moved += p.stat().st_size
                except OSError:
                    pass
                files_moved += 1
        if i % 25 == 0 or i == len(module_ids):
            print(f"[snapshot] {i}/{len(module_ids)} modules moved "
                  f"({files_moved:,} files, {bytes_moved/1024/1024:.1f} MB)")

    return modules_moved, files_moved, bytes_moved


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--hub-data-dir", required=True,
                        help="Hub --data-dir. Place on the same volume as --snapshot-dir so the snapshot step is a rename per module instead of a cross-volume byte copy.")
    parser.add_argument("--snapshot-dir", required=True,
                        help="Destination for per-module server-state trees.")
    parser.add_argument("--port", type=int, default=8558,
                        help="Hub HTTP port (default: 8558)")
    parser.add_argument("--module-count", type=int, default=1000,
                        help="Number of modules to snapshot (default: 300)")
    parser.add_argument("--workers", type=int, default=50,
                        help="Concurrent HTTP workers (default: 50)")
    parser.add_argument("--poll-timeout", type=float, default=1800.0,
                        help="Max seconds to wait for provision or hibernate to finish (default: 1800)")
    parser.add_argument("--zenserver-dir",
                        help="Directory containing zenserver executable (auto-detected by default)")
    args = parser.parse_args()

    hub_data_dir = Path(args.hub_data_dir).resolve()
    snapshot_dir = Path(args.snapshot_dir).resolve()

    # Safety: snapshot-dir is the preserved output. It must not overlap the
    # hub data-dir in either direction so nothing the hub writes can clobber
    # the preserved state.
    if (snapshot_dir == hub_data_dir
            or snapshot_dir in hub_data_dir.parents
            or hub_data_dir in snapshot_dir.parents):
        sys.exit(f"[setup] snapshot-dir ({snapshot_dir}) and hub-data-dir ({hub_data_dir}) must be disjoint")

    s3_uri = os.environ.get("S3_URI", _DEFAULT_S3_URI)
    aws_profile = os.environ.get("AWS_PROFILE", _DEFAULT_AWS_PROFILE)
    aws_region = os.environ.get("AWS_REGION", _DEFAULT_AWS_REGION)
    if not s3_uri:
        sys.exit("[setup] S3 URI not set. Set ZEN_PERF_S3_URI (or S3_URI) to a bucket like s3://your-bucket/")
    if not aws_profile:
        sys.exit("[setup] AWS profile not set. Set ZEN_PERF_AWS_PROFILE (or AWS_PROFILE) to your SSO profile name")

    hub_log = hub_data_dir / "hub.log"

    zenserver_exe = _find_zenserver(args.zenserver_dir)
    zen_exe = _find_zen(zenserver_exe)
    zenserver_mode = "release" if "release" in zenserver_exe.parts else ("debug" if "debug" in zenserver_exe.parts else "?")
    print(f"[setup] build mode:   {zenserver_mode}")
    print(f"[setup] zenserver:    {zenserver_exe}")
    print(f"[setup] zen cli:      {zen_exe}")
    print(f"[setup] S3 URI:       {s3_uri}")
    print(f"[setup] profile:      {aws_profile}")
    print(f"[setup] hub-data-dir: {hub_data_dir}")
    print(f"[setup] snapshot-dir: {snapshot_dir}")

    session, frozen = _get_session(aws_profile)
    print(f"[aws] credentials resolved (key prefix {frozen.access_key[:6]}..., session-token={'yes' if frozen.token else 'no'})")

    bucket, prefix = _parse_s3_uri(s3_uri)
    print(f"[s3] listing folders under bucket='{bucket}' prefix='{prefix}'...")
    module_ids = _list_module_ids(session, bucket, prefix, aws_region, args.module_count)
    print(f"[s3] selected {len(module_ids)} module folders (UUID-shaped)")
    if len(module_ids) < args.module_count:
        print(f"[s3] WARNING: asked for {args.module_count} modules, only {len(module_ids)} matched the UUID filter")

    # _list_module_ids returns modules grouped by size bucket (small first, then
    # medium, then large). Provisioning in that order would create a size-sorted
    # fan-out and skew load: all small first (cheap, fast hydrations) then all
    # large (heavy, slow). Shuffle so provision/hibernate/copy hit a mixed-size
    # stream that better resembles real workload distribution. Fixed seed keeps
    # runs reproducible across reseeds with the same bucket contents.
    import random
    random.Random(0xC0FFEE).shuffle(module_ids)

    if not module_ids:
        sys.exit("[s3] no module folders found, aborting")

    aws_env = {
        "AWS_ACCESS_KEY_ID": frozen.access_key,
        "AWS_SECRET_ACCESS_KEY": frozen.secret_key,
    }
    if frozen.token:
        aws_env["AWS_SESSION_TOKEN"] = frozen.token

    hub_data_dir.mkdir(parents=True, exist_ok=True)
    config_path = hub_data_dir / "hydration_config.json"
    config_path.write_text(
        json.dumps({
            "type": "s3",
            "settings": {"uri": s3_uri, "region": aws_region},
        }),
        encoding="ascii",
    )
    hub_extra_args = [
        f"--hub-hydration-target-config={config_path}",
        "--hub-enable-dehydration=false",
    ]

    hub_proc: Optional[subprocess.Popen] = None
    hub_log_handle = None
    exit_code = 0

    try:
        hub_instance_limit = max(args.module_count + 10, 500)
        hub_proc, hub_log_handle = _start_hub(
            zenserver_exe, hub_data_dir, args.port, hub_log,
            hub_instance_limit, hub_extra_args, aws_env,
        )
        _wait_for_hub(hub_proc, args.port)

        t_start = time.monotonic()

        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            # --- Provision ---
            print(f"[provision] firing {len(module_ids)} provision requests...")
            t0 = time.monotonic()
            accepted, failures = _fan_out_post(pool, args.port, module_ids, "provision")
            print(f"[provision] accepted={len(accepted)}, failed={len(failures)} (fan-out {time.monotonic()-t0:.1f}s)")
            for mid, status, body in failures[:10]:
                print(f"[provision] FAILED {mid}: status={status} body={body}")
            if not accepted:
                sys.exit("[provision] nothing accepted, aborting")

            hydration_warned: set[str] = set()
            stuck, failed, last_states = _wait_for_state(
                args.port, accepted, "provisioned", args.poll_timeout, "provision",
                hub_log=hub_log, hydration_warned=hydration_warned,
            )
            prov_done = len(accepted) - len(stuck) - len(failed)
            print(f"[provision] complete: {prov_done}/{len(accepted)} provisioned, {len(failed)} failed, {len(stuck)} stuck, "
                  f"{len(hydration_warned)} hydration-failed ({time.monotonic()-t0:.1f}s)")
            if failed:
                for mid in failed[:10]:
                    print(f"[provision]   FAILED {mid}: last state='{last_states.get(mid, '')}'")
                exit_code = 1
            if stuck:
                for mid in stuck[:10]:
                    print(f"[provision]   STUCK {mid}: last state='{last_states.get(mid, '')}'")
                exit_code = 1
            accepted = [m for m in accepted if m not in set(stuck) and m not in set(failed)]

            if not accepted:
                sys.exit("[provision] nothing successfully provisioned, aborting")

            # --- Hibernate ---
            print(f"[hibernate] firing {len(accepted)} hibernate requests...")
            t0 = time.monotonic()
            hib_accepted, hib_failures = _fan_out_post(pool, args.port, accepted, "hibernate")
            print(f"[hibernate] accepted={len(hib_accepted)}, failed={len(hib_failures)} (fan-out {time.monotonic()-t0:.1f}s)")
            for mid, status, body in hib_failures[:10]:
                print(f"[hibernate] FAILED {mid}: status={status} body={body}")
                exit_code = 1

            stuck_hib, failed_hib, last_states_hib = _wait_for_state(args.port, hib_accepted, "hibernated", args.poll_timeout, "hibernate")
            hib_done = len(hib_accepted) - len(stuck_hib) - len(failed_hib)
            print(f"[hibernate] complete: {hib_done}/{len(hib_accepted)} hibernated, {len(failed_hib)} failed, {len(stuck_hib)} stuck "
                  f"({time.monotonic()-t0:.1f}s)")
            if failed_hib or stuck_hib:
                exit_code = 1
                for mid in (failed_hib + stuck_hib)[:10]:
                    print(f"[hibernate]   not-hibernated {mid}: last state='{last_states_hib.get(mid, '')}'")

        # --- Snapshot per-module state out of the hub data dir while the hub
        #     is still running. All instances are hibernated (no writers),
        #     watchdog-hibernated-timeout is 86400s (no auto-deprovision),
        #     hub is only touching its own metadata outside servers/<mid>/.
        #     When --hub-data-dir and --snapshot-dir share a volume the per-
        #     module trees are renamed (O(1)); cross-volume falls back to a
        #     byte copy. The hub data dir is wiped on the next run regardless.
        snapshot_src = hub_data_dir / "servers"
        to_snapshot = [m for m in hib_accepted if m not in set(stuck_hib) and m not in set(failed_hib)]
        print(f"[snapshot] moving {len(to_snapshot)} module trees from {snapshot_src} -> {snapshot_dir}")
        t0 = time.monotonic()
        modules_moved, files_moved, bytes_moved = _copy_snapshot(snapshot_src, snapshot_dir, to_snapshot)
        print(f"[snapshot] moved {modules_moved} modules, {files_moved:,} files, {bytes_moved/1024/1024:.1f} MB "
              f"({time.monotonic()-t0:.1f}s)")

        if modules_moved < len(to_snapshot):
            print(f"[snapshot] WARNING: only {modules_moved}/{len(to_snapshot)} trees moved")
            exit_code = 1

        # --- Graceful hub shutdown via 'zen down' ---
        _zen_down_hub(zen_exe, hub_proc)
        hub_proc = None
        if hub_log_handle is not None:
            hub_log_handle.close()
            hub_log_handle = None

        print(f"[summary] stage A total elapsed: {time.monotonic()-t_start:.1f}s, exit={exit_code}")
        print(f"[summary] snapshot available at: {snapshot_dir}")

    finally:
        if hub_proc is not None and hub_proc.poll() is None:
            # Fallback: something aborted before the 'zen down' path ran.
            print("[hub] force-terminating (zen down was not invoked)")
            hub_proc.terminate()
            try:
                hub_proc.wait(timeout=60)
            except subprocess.TimeoutExpired:
                hub_proc.kill()
                hub_proc.wait()
        if hub_log_handle is not None:
            hub_log_handle.close()

    return exit_code


if __name__ == "__main__":
    sys.exit(main())