Stream model reserve to Scaleway without local weight staging

Assistant: codex
Assistant-Model: gpt-6-astra
Assistant-Session: 01a09cbd-43c1-79f3-809e-1ee97b40b64d

scripts/collect_model.py

@@ -1,27 +1,15 @@
 #!/usr/bin/env python3
-"""Collect one HF model into local staging, optionally upload to Scaleway.
+"""Collect an HF model: diskless S3 streaming or explicitly opted-in local cache.
 
-Usage:
-  python3 scripts/collect_model.py \\
-    --repo-id Qwen/Qwen3-8B --local-name Qwen__Qwen3-8B --revision main
-
-  python3 scripts/collect_model.py \\
-    --repo-id deepseek-ai/DeepSeek-V4-Flash-0731 \\
-    --local-name deepseek-ai__DeepSeek-V4-Flash-0731 \\
-    --tier s \\
+Remote reserve (default when --s3-bucket is supplied; no local weight files):
+  python3 scripts/collect_model.py \
+    --repo-id deepseek-ai/DeepSeek-V4-Flash-0731 \
+    --local-name deepseek-ai__DeepSeek-V4-Flash-0731 --tier s \
     --s3-bucket railiance-fi-open-weight-reserve
 
-Layout (local staging, still required for huggingface_hub):
-  {base}/models/{local_name}/{revision}/   # R / W
-  {base}/strategic/{local_name}/{revision}/  # S
-  {base}/staging/...
-
-S3 prefix (SoT once bucket is live):
-  s3://{bucket}/models|strategic/{local_name}/{revision}/
-  s3://{bucket}/manifests/{local_name}/{revision}/MANIFEST.json
-
-Credentials: AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY, or SCW_ACCESS_KEY +
-SCW_SECRET_KEY. Never commit keys.
+Use --dry-run to inspect the immutable source revision and transfer size.
+Use --local-download without --s3-bucket only for an intentional local cache.
+See docs/streaming-reserve.md for the resource-limited remote worker.
 """
 from __future__ import annotations
 
@@ -36,8 +24,6 @@ from pathlib import Path
 # Prefer sequential downloads on WSL+drvfs (xet parallel can OOM)
 os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
 
-from huggingface_hub import snapshot_download  # noqa: E402
-
 DEFAULT_ALLOW = (
     "*.safetensors",
     "*.bin",
@@ -112,8 +98,22 @@ def main() -> int:
         default=os.environ.get("FI_S3_REGION", "nl-ams"),
     )
     ap.add_argument("--token", default=os.environ.get("HF_TOKEN") or None)
+    ap.add_argument("--dry-run", action="store_true", help="list pinned S3 transfer without uploading")
+    ap.add_argument("--part-mib", type=int, default=64, help="streaming S3 part buffer (5–128 MiB)")
+    ap.add_argument("--local-download", action="store_true", help="explicitly permit a full local snapshot (no S3)")
     args = ap.parse_args()
 
+    if args.s3_bucket:
+        if args.local_download:
+            ap.error("--local-download cannot be combined with --s3-bucket")
+        from stream_model import collect
+        return collect(args, DEFAULT_ALLOW, IGNORE)
+    if not args.local_download:
+        ap.error("choose --s3-bucket (diskless streaming) or explicitly --local-download")
+    if args.dry_run:
+        ap.error("--dry-run currently requires --s3-bucket")
+    from huggingface_hub import snapshot_download
+
     base = Path(args.base)
     dest_kind = "strategic" if args.tier == "s" else "models"
     stage = base / "staging" / args.local_name / args.revision
@@ -179,97 +179,9 @@ def main() -> int:
     (final / "MANIFEST.json").write_text(json.dumps(manifest, indent=2) + "\n")
     result = {"ok": True, "total_bytes": total, "path": str(final), "tier": args.tier}
 
-    if args.s3_bucket:
-        prefix = f"{dest_kind}/{args.local_name}/{args.revision}"
-        storage_class = "GLACIER" if args.tier == "s" else "ONEZONE_IA"
-        uploaded = upload_tree(
-            final,
-            bucket=args.s3_bucket,
-            prefix=prefix,
-            endpoint=args.s3_endpoint,
-            region=args.s3_region,
-            storage_class=storage_class,
-        )
-        man_key = f"manifests/{args.local_name}/{args.revision}/MANIFEST.json"
-        upload_file(
-            final / "MANIFEST.json",
-            bucket=args.s3_bucket,
-            key=man_key,
-            endpoint=args.s3_endpoint,
-            region=args.s3_region,
-            storage_class="STANDARD",
-        )
-        result["s3"] = {
-            "bucket": args.s3_bucket,
-            "prefix": f"s3://{args.s3_bucket}/{prefix}/",
-            "manifest": f"s3://{args.s3_bucket}/{man_key}",
-            "storage_class": storage_class,
-            "objects": uploaded,
-        }
-
     print(json.dumps(result, indent=2))
     return 0
 
 
-def _s3_client(endpoint: str, region: str):
-    try:
-        import boto3
-    except ImportError as exc:
-        raise SystemExit("boto3 required for --s3-bucket (pip install boto3)") from exc
-    access = os.environ.get("AWS_ACCESS_KEY_ID") or os.environ.get("SCW_ACCESS_KEY")
-    secret = os.environ.get("AWS_SECRET_ACCESS_KEY") or os.environ.get("SCW_SECRET_KEY")
-    if not access or not secret:
-        raise SystemExit("set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY or SCW_ACCESS_KEY/SCW_SECRET_KEY")
-    return boto3.client(
-        "s3",
-        region_name=region,
-        endpoint_url=endpoint,
-        aws_access_key_id=access,
-        aws_secret_access_key=secret,
-    )
-
-
-def upload_file(
-    path: Path,
-    *,
-    bucket: str,
-    key: str,
-    endpoint: str,
-    region: str,
-    storage_class: str,
-) -> None:
-    client = _s3_client(endpoint, region)
-    extra = {"StorageClass": storage_class} if storage_class else {}
-    client.upload_file(str(path), bucket, key, ExtraArgs=extra)
-
-
-def upload_tree(
-    root: Path,
-    *,
-    bucket: str,
-    prefix: str,
-    endpoint: str,
-    region: str,
-    storage_class: str,
-) -> int:
-    n = 0
-    for p in sorted(root.rglob("*")):
-        if not p.is_file() or p.name.startswith("."):
-            continue
-        rel = p.relative_to(root).as_posix()
-        key = f"{prefix.rstrip('/')}/{rel}"
-        upload_file(
-            p,
-            bucket=bucket,
-            key=key,
-            endpoint=endpoint,
-            region=region,
-            storage_class=storage_class,
-        )
-        n += 1
-        print(f"uploaded {key}", file=sys.stderr)
-    return n
-
-
 if __name__ == "__main__":
     sys.exit(main())