freedom-intelligence/scripts/stream_model.py

"""Diskless HF → S3 transfer. One file/part at a time; durable per-file receipts.

Run through collect_model.py. A single worker must own each destination prefix.
No HF download/cache functions are used. Interrupted files restart from byte zero;
completed files resume from receipts after checking the remote object identity.
"""
from __future__ import annotations

import base64
from datetime import datetime, timezone
import fnmatch
import hashlib
import json
import math
import os
import re
import signal
import sys
import time
from pathlib import Path


def md5(data):
    return base64.b64encode(hashlib.md5(data).digest()).decode()


def client(endpoint, region):
    import boto3
    from botocore.config import Config

    credentials = {}
    if os.environ.get("FI_S3_CREDENTIAL_FILE"):
        credentials = json.loads(Path(os.environ["FI_S3_CREDENTIAL_FILE"]).read_text())
    access = (credentials.get("ACCESS_KEY") or os.environ.get("AWS_ACCESS_KEY_ID")
              or os.environ.get("SCW_ACCESS_KEY"))
    secret = (credentials.get("SECRET_KEY") or os.environ.get("AWS_SECRET_ACCESS_KEY")
              or os.environ.get("SCW_SECRET_KEY"))
    if not access or not secret:
        raise RuntimeError("S3 credentials missing")
    return boto3.client(
        "s3", endpoint_url=endpoint, region_name=region,
        aws_access_key_id=access, aws_secret_access_key=secret,
        aws_session_token=credentials.get("SESSION_TOKEN") or os.environ.get("AWS_SESSION_TOKEN"),
        config=Config(connect_timeout=20, read_timeout=180,
                      retries={"max_attempts": 5, "mode": "standard"},
                      request_checksum_calculation="when_required",
                      response_checksum_validation="when_required"),
    )


def put_json(s3, bucket, key, value):
    data = (json.dumps(value, indent=2) + "\n").encode()
    result = s3.put_object(Bucket=bucket, Key=key, Body=data, ContentMD5=md5(data),
                           ContentType="application/json", StorageClass="STANDARD")
    if result["ETag"].strip('"') != hashlib.md5(data).hexdigest():
        raise ValueError(f"metadata object checksum mismatch: {key}")
    if optional_json(s3, bucket, key) != value:
        raise ValueError(f"metadata object readback mismatch: {key}")


def optional_json(s3, bucket, key):
    from botocore.exceptions import ClientError
    try:
        response = s3.get_object(Bucket=bucket, Key=key)
    except ClientError as exc:
        if exc.response["Error"]["Code"] in {"NoSuchKey", "404"}:
            return None
        raise
    with response["Body"] as body:
        return json.load(body)


def source_file(sibling):
    lfs = sibling.lfs
    if sibling.size is None:
        raise ValueError(f"no source size for {sibling.rfilename}")
    digest = lfs.sha256 if lfs else sibling.blob_id
    algorithm = "sha256" if lfs else "git-sha1"
    if not digest or not re.fullmatch(r"[0-9a-f]{64}" if lfs else r"[0-9a-f]{40}", digest):
        raise ValueError(f"no usable upstream digest for {sibling.rfilename}")
    return {"path": sibling.rfilename, "bytes": sibling.size,
            "source_digest": digest, "source_algorithm": algorithm}


def stream_object(s3, bucket, key, body, spec, storage_class, part_size, metadata):
    """Verify source hash and returned part/composite MD5 ETags before receipt.

    Scaleway's live API accepted incorrect Content-MD5 (2026-09-14), so sending
    that header alone is insufficient. Fail closed if ETags are not MD5-shaped.
    """
    if math.ceil(spec["bytes"] / part_size) > 10000:
        raise ValueError("file exceeds 10,000 parts; increase --part-mib")
    sha = hashlib.sha256()
    git_sha = hashlib.sha1(f'blob {spec["bytes"]}\0'.encode())
    upload_id = None
    total = 0
    parts = []
    part_digests = []
    try:
        if spec["bytes"]:
            upload_id = s3.create_multipart_upload(
                Bucket=bucket, Key=key, StorageClass=storage_class, Metadata=metadata,
            )["UploadId"]
        while True:
            # read(n) from urllib3's HTTPResponse fills n bytes except at EOF.
            block = body.read(part_size)
            if not block:
                break
            total += len(block)
            if total > spec["bytes"]:
                raise ValueError("source exceeded declared size")
            sha.update(block)
            git_sha.update(block)
            result = s3.upload_part(
                Bucket=bucket, Key=key, UploadId=upload_id,
                PartNumber=len(parts) + 1, Body=block, ContentMD5=md5(block),
            )
            part_digest = hashlib.md5(block).digest()
            if result['ETag'].strip('"') != part_digest.hex():
                raise ValueError(f"destination part checksum mismatch: {key}")
            part_digests.append(part_digest)
            parts.append({"PartNumber": len(parts) + 1, "ETag": result["ETag"]})
            print(f"part {key} {total}/{spec['bytes']}", flush=True)
            del block
        actual = sha.hexdigest() if spec["source_algorithm"] == "sha256" else git_sha.hexdigest()
        if total != spec["bytes"] or actual != spec["source_digest"]:
            raise ValueError(f"source size/checksum mismatch: {key}")
        if upload_id:
            expected_etag = hashlib.md5(b"".join(part_digests)).hexdigest() + f"-{len(parts)}"
            result = s3.complete_multipart_upload(
                Bucket=bucket, Key=key, UploadId=upload_id, MultipartUpload={"Parts": parts},
            )
            upload_id = None
        else:
            expected_etag = hashlib.md5(b"").hexdigest()
            result = s3.put_object(Bucket=bucket, Key=key, Body=b"", ContentMD5=md5(b""),
                                   StorageClass=storage_class, Metadata=metadata)
        head = s3.head_object(Bucket=bucket, Key=key)
        if (head["ContentLength"] != total or head["Metadata"] != metadata
                or head["ETag"] != result["ETag"]
                or head["ETag"].strip('"') != expected_etag):
            raise ValueError(f"destination identity mismatch: {key}")
        return dict(spec, sha256=sha.hexdigest(), key=key, etag=head["ETag"],
                    version_id=head.get("VersionId"), storage_class=storage_class,
                    verification="upstream-digest+part-md5-etags+composite-etag+head; no restore readback")
    finally:
        if upload_id:
            s3.abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)


def receipt_matches(s3, bucket, key, receipt, spec, metadata, storage_class):
    from botocore.exceptions import ClientError
    if not receipt or any(receipt.get(k) != v for k, v in spec.items()):
        return False
    if receipt.get("key") != key or receipt.get("storage_class") != storage_class:
        return False
    if not re.fullmatch(r"[0-9a-f]{64}", receipt.get("sha256", "")):
        return False
    try:
        head = s3.head_object(Bucket=bucket, Key=key)
    except ClientError as exc:
        if exc.response["Error"]["Code"] in {"NoSuchKey", "404", "NotFound"}:
            return False
        raise
    return (head["ContentLength"] == spec["bytes"]
            and head["ETag"] == receipt.get("etag")
            and head.get("VersionId") == receipt.get("version_id")
            and head["Metadata"] == metadata
            and head.get("StorageClass", "STANDARD") == storage_class)


def abort_orphans(s3, bucket, prefix):
    # Only this collector's immutable model prefix; never bucket-wide cleanup.
    paginator = s3.get_paginator("list_multipart_uploads")
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix + "/"):
        for upload in page.get("Uploads", []):
            s3.abort_multipart_upload(Bucket=bucket, Key=upload["Key"], UploadId=upload["UploadId"])


def collect(args, allow, ignore):
    import requests
    from huggingface_hub import HfApi, hf_hub_url

    if not 5 <= args.part_mib <= 128:
        raise ValueError("--part-mib must be between 5 and 128")
    if not re.fullmatch(r"[A-Za-z0-9_.-]+", args.local_name) or args.local_name in {".", ".."}:
        raise ValueError("--local-name must be one safe path component")
    info = HfApi(token=args.token).model_info(args.repo_id, revision=args.revision, files_metadata=True)
    if not re.fullmatch(r"[0-9a-f]{40}", info.sha):
        raise ValueError("source did not resolve to an immutable commit")
    files = [source_file(f) for f in info.siblings
             if any(fnmatch.fnmatchcase(f.rfilename, p) for p in allow)
             and not any(fnmatch.fnmatchcase(f.rfilename, p) for p in ignore)]
    files.sort(key=lambda f: f["path"])
    if not files:
        raise ValueError("no matching source files")
    part_size = args.part_mib * 1024 * 1024
    if any(math.ceil(f["bytes"] / part_size) > 10000 for f in files):
        raise ValueError("file exceeds 10,000 parts; increase --part-mib")
    kind = "strategic" if args.tier == "s" else "models"
    prefix = f"{kind}/{args.local_name}/{info.sha}"
    manifest_prefix = f"manifests/{args.local_name}/{info.sha}"
    plan = {"repo_id": args.repo_id, "revision": info.sha, "requested_revision": args.revision,
            "storage_path": f"s3://{args.s3_bucket}/{prefix}/", "local_name": args.local_name,
            "total_bytes": sum(f["bytes"] for f in files), "files": len(files),
            "part_mib": args.part_mib, "local_weight_bytes": 0}
    print(json.dumps(plan), flush=True)
    if args.dry_run:
        return 0
    s3 = client(args.s3_endpoint, args.s3_region)
    s3.head_bucket(Bucket=args.s3_bucket)
    abort_orphans(s3, args.s3_bucket, prefix)
    storage_class = "GLACIER" if args.tier == "s" else "ONEZONE_IA"
    artifacts = []
    with requests.Session() as session:
        for spec in files:
            key = f"{prefix}/{spec['path']}"
            receipt_key = f"{manifest_prefix}/receipts/{hashlib.sha256(spec['path'].encode()).hexdigest()}.json"
            metadata = {"fi-revision": info.sha, "fi-source-digest": spec["source_digest"],
                        "fi-source-algorithm": spec["source_algorithm"]}
            receipt = optional_json(s3, args.s3_bucket, receipt_key)
            if receipt_matches(s3, args.s3_bucket, key, receipt, spec, metadata, storage_class):
                print(f"resume verified object {key}", flush=True)
                artifacts.append(receipt)
                continue
            for attempt in range(3):
                try:
                    headers = {"Accept-Encoding": "identity"}
                    if args.token:
                        headers["Authorization"] = f"Bearer {args.token}"
                    # requests strips Authorization on cross-host CDN redirects.
                    with session.get(hf_hub_url(args.repo_id, spec["path"], revision=info.sha),
                                     headers=headers, stream=True, timeout=(20, 180)) as response:
                        response.raise_for_status()
                        if response.status_code != 200:
                            raise ValueError("expected a full source response")
                        if response.headers.get("Content-Encoding", "identity") != "identity":
                            raise ValueError("unexpected source content encoding")
                        receipt = stream_object(s3, args.s3_bucket, key, response.raw, spec,
                                                storage_class, part_size, metadata)
                    put_json(s3, args.s3_bucket, receipt_key, receipt)
                    artifacts.append(receipt)
                    break
                except Exception as exc:
                    # Never print signed source URLs or credentials in exception text.
                    print(f"retry {spec['path']} attempt={attempt + 1} error={type(exc).__name__}",
                          file=sys.stderr, flush=True)
                    if attempt == 2:
                        raise RuntimeError(f"transfer failed: {spec['path']} ({type(exc).__name__})") from None
                    time.sleep(5 * (attempt + 1))
    manifest = dict(plan, artifacts=artifacts, completed_at=datetime.now(timezone.utc).isoformat(),
                    complete=True, transport="http-stream-to-s3-multipart")
    manifest_key = f"{manifest_prefix}/MANIFEST.json"
    put_json(s3, args.s3_bucket, manifest_key, manifest)
    print(json.dumps({"ok": True, "manifest": f"s3://{args.s3_bucket}/{manifest_key}"}), flush=True)
    return 0


def terminate(signum, frame):
    raise SystemExit(128 + signum)


# SIGTERM unwinds the current multipart upload; SIGKILL is cleaned on next run.
signal.signal(signal.SIGTERM, terminate)