scripts/omni/parallel_chunk_transfer.py

#!/usr/bin/env python3
"""Transfer large dataset files over several SSH/SCP streams.

This utility is intended for relay situations where a single rsync/ssh stream is
slow. Files above a configurable threshold are split locally, copied as chunks in
parallel, reassembled on the destination host, and size-validated before the
final file is moved into place.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import fnmatch
import json
import os
import shlex
import shutil
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--src-root", type=Path, required=True)
    parser.add_argument("--dst-host", required=True, help="Destination SSH host, e.g. cy@host")
    parser.add_argument("--dst-root", required=True)
    parser.add_argument("--ssh-key", type=Path, required=True)
    parser.add_argument("--ssh-extra", default="-o BatchMode=yes -o StrictHostKeyChecking=accept-new")
    parser.add_argument("--chunk-size-mib", type=int, default=128)
    parser.add_argument("--parallel", type=int, default=8)
    parser.add_argument("--split-threshold-mib", type=int, default=256)
    parser.add_argument("--local-temp-root", type=Path, default=None)
    parser.add_argument("--remote-temp-root", default="/tmp/xperience10m_chunk_transfer")
    parser.add_argument("--exclude", action="append", default=["visualization.rrd", ".cache/**", "**/.cache/**", ".chunk_transfer_tmp/**"])
    parser.add_argument("--max-files", type=int, default=0, help="0 means all files.")
    parser.add_argument("--progress-jsonl", type=Path, default=None)
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def human_bytes(num: int | float) -> str:
    value = float(num)
    for unit in ["B", "KiB", "MiB", "GiB", "TiB"]:
        if abs(value) < 1024.0 or unit == "TiB":
            return f"{value:.2f} {unit}"
        value /= 1024.0
    return f"{value:.2f} TiB"


def log_event(path: Path | None, payload: dict) -> None:
    payload = {"time": utc_now(), **payload}
    line = json.dumps(payload, sort_keys=True)
    print(line, flush=True)
    if path is not None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("a", encoding="utf-8") as handle:
            handle.write(line + "\n")


def run(cmd: list[str], dry_run: bool = False) -> None:
    if dry_run:
        print("DRY-RUN", " ".join(shlex.quote(part) for part in cmd), flush=True)
        return
    subprocess.run(cmd, check=True)


def ssh_base(args: argparse.Namespace) -> list[str]:
    return ["ssh", "-i", str(args.ssh_key), *shlex.split(args.ssh_extra), args.dst_host]


def scp_base(args: argparse.Namespace) -> list[str]:
    return ["scp", "-q", "-i", str(args.ssh_key), *shlex.split(args.ssh_extra)]


def excluded(rel_path: str, patterns: list[str]) -> bool:
    name = Path(rel_path).name
    if rel_path.startswith(".cache/") or "/.cache/" in rel_path:
        return True
    if ".chunk_transfer_tmp/" in rel_path or rel_path.startswith(".chunk_transfer_tmp/"):
        return True
    return any(fnmatch.fnmatch(rel_path, pattern) or fnmatch.fnmatch(name, pattern) for pattern in patterns)


def iter_files(root: Path, patterns: list[str]) -> list[Path]:
    files: list[Path] = []
    for path in root.rglob("*"):
        if not path.is_file():
            continue
        rel = path.relative_to(root).as_posix()
        if excluded(rel, patterns):
            continue
        files.append(path)
    return sorted(files)


def remote_mkdir(args: argparse.Namespace, path: str) -> None:
    run([*ssh_base(args), "mkdir", "-p", path], args.dry_run)


def remote_size(args: argparse.Namespace, path: str) -> int | None:
    code = "import os,sys; p=sys.argv[1]; print(os.path.getsize(p) if os.path.exists(p) else -1)"
    remote_cmd = f"python3 -c {shlex.quote(code)} {shlex.quote(path)}"
    cmd = [*ssh_base(args), remote_cmd]
    if args.dry_run:
        return None
    out = subprocess.check_output(cmd, text=True).strip()
    size = int(out)
    return None if size < 0 else size


def copy_one(args: argparse.Namespace, src: Path, remote_path: str) -> None:
    remote_mkdir(args, shlex.quote(str(Path(remote_path).parent)))
    run([*scp_base(args), str(src), f"{args.dst_host}:{remote_path}"], args.dry_run)


def split_file(src: Path, chunk_dir: Path, chunk_size: int) -> list[Path]:
    chunk_dir.mkdir(parents=True, exist_ok=True)
    chunks: list[Path] = []
    with src.open("rb") as handle:
        index = 0
        while True:
            data = handle.read(chunk_size)
            if not data:
                break
            chunk = chunk_dir / f"part_{index:05d}"
            chunk.write_bytes(data)
            chunks.append(chunk)
            index += 1
    return chunks


def transfer_chunk(args: argparse.Namespace, chunk: Path, remote_chunk_dir: str) -> None:
    run([*scp_base(args), str(chunk), f"{args.dst_host}:{remote_chunk_dir}/{chunk.name}"], args.dry_run)


def assemble_remote(args: argparse.Namespace, remote_chunk_dir: str, chunk_names: list[str], remote_path: str, expected_size: int) -> None:
    code = (
        "import os, shutil, sys\n"
        "chunk_dir, dest, expected, *names = sys.argv[1:]\n"
        "expected = int(expected)\n"
        "os.makedirs(os.path.dirname(dest), exist_ok=True)\n"
        "partial = dest + '.partial_chunked'\n"
        "with open(partial, 'wb') as out:\n"
        "    for name in names:\n"
        "        with open(os.path.join(chunk_dir, name), 'rb') as inp:\n"
        "            shutil.copyfileobj(inp, out, 1024 * 1024)\n"
        "actual = os.path.getsize(partial)\n"
        "if actual != expected:\n"
        "    raise SystemExit(f'size mismatch for {dest}: expected {expected}, got {actual}')\n"
        "os.replace(partial, dest)\n"
    )
    argv = [remote_chunk_dir, remote_path, str(expected_size), *chunk_names]
    remote_cmd = " ".join(["python3", "-c", shlex.quote(code), *(shlex.quote(value) for value in argv)])
    run([*ssh_base(args), remote_cmd], args.dry_run)


def transfer_split(args: argparse.Namespace, src: Path, remote_path: str, local_temp: Path, remote_temp: str) -> None:
    size = src.stat().st_size
    if remote_size(args, remote_path) == size:
        log_event(args.progress_jsonl, {"event": "skip_existing", "path": str(src), "bytes": size})
        return

    if local_temp.exists():
        shutil.rmtree(local_temp)
    chunks = split_file(src, local_temp, args.chunk_size_mib * 1024 * 1024)
    remote_mkdir(args, remote_temp)
    start = time.time()
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.parallel) as pool:
        futures = [pool.submit(transfer_chunk, args, chunk, remote_temp) for chunk in chunks]
        for future in concurrent.futures.as_completed(futures):
            future.result()
    assemble_remote(args, remote_temp, [chunk.name for chunk in chunks], remote_path, size)
    elapsed = time.time() - start
    rate = size / elapsed if elapsed else 0
    log_event(
        args.progress_jsonl,
        {
            "event": "split_transfer_done",
            "path": str(src),
            "bytes": size,
            "human": human_bytes(size),
            "chunks": len(chunks),
            "seconds": round(elapsed, 3),
            "rate_human_per_s": f"{human_bytes(rate)}/s",
        },
    )


def main() -> int:
    args = parse_args()
    args.src_root = args.src_root.expanduser().resolve()
    args.ssh_key = args.ssh_key.expanduser().resolve()
    if args.local_temp_root is None:
        args.local_temp_root = args.src_root.parent / f".{args.src_root.name}_chunk_transfer_tmp"
    else:
        args.local_temp_root = args.local_temp_root.expanduser().resolve()

    files = iter_files(args.src_root, args.exclude)
    if args.max_files:
        files = files[: args.max_files]
    threshold = args.split_threshold_mib * 1024 * 1024
    log_event(
        args.progress_jsonl,
        {
            "event": "start",
            "src_root": str(args.src_root),
            "dst_host": args.dst_host,
            "dst_root": args.dst_root,
            "files": len(files),
            "parallel": args.parallel,
            "chunk_size_mib": args.chunk_size_mib,
        },
    )

    for index, src in enumerate(files):
        rel = src.relative_to(args.src_root).as_posix()
        remote_path = f"{args.dst_root.rstrip('/')}/{rel}"
        size = src.stat().st_size
        log_event(args.progress_jsonl, {"event": "file_start", "index": index, "path": rel, "bytes": size, "human": human_bytes(size)})
        if size >= threshold:
            safe_name = rel.replace("/", "__")
            transfer_split(
                args,
                src,
                remote_path,
                args.local_temp_root / safe_name,
                f"{args.remote_temp_root.rstrip('/')}/{safe_name}",
            )
        else:
            start = time.time()
            copy_one(args, src, remote_path)
            elapsed = time.time() - start
            log_event(
                args.progress_jsonl,
                {
                    "event": "direct_transfer_done",
                    "path": rel,
                    "bytes": size,
                    "seconds": round(elapsed, 3),
                    "rate_human_per_s": f"{human_bytes(size / elapsed if elapsed else 0)}/s",
                },
            )

    log_event(args.progress_jsonl, {"event": "done"})
    return 0


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