NIM air-gap cache processing script¶

#!/usr/bin/env python3
"""
Uploads a local NIM cache directly to S3/MinIO without modifying the filesystem.

## Steps {: #steps }
1. Walk the NIM cache, resolving symlinks to real file paths.
2. Compute the URL-encoded NIM key for each file.
3. Upload directly to S3/MinIO using the computed key.

No files are moved, copied, or deleted on disk.

## Arguments {: #arguments }
--cache-dir            Path to the local NIM cache.        (required)
--bucket-name          Destination bucket name.            (required)
--endpoint-url         S3/MinIO endpoint URL.              (optional, omit for AWS S3)
--access-key           S3/MinIO access key.                (optional, falls back to env/IAM)
--secret-key           S3/MinIO secret key.                (optional, falls back to env/IAM)
--region               AWS region. Defaults to us-east-1.  (optional)
--insecure             Skip SSL certificate verification.  (optional)
--dry-run              Print keys without uploading.       (optional)

## Usage examples {: #usage-examples }
python process_nim_cache.py \\
    --cache-dir /opt/nim/.cache \\
    --bucket-name nim-models \\
    --endpoint-url https://minio.internal:9000 \\
    --access-key miniorootuser \\
    --secret-key secret \\
    --insecure

python process_nim_cache.py \\
    --cache-dir /opt/nim/.cache \\
    --bucket-name nim-models \\
    --endpoint-url https://minio.internal:9000 \\
    --access-key miniorootuser \\
    --secret-key secret \\
    --dry-run
"""
import os, argparse, urllib3, urllib.parse
from pathlib import Path


def iter_keys(cache_dir):
    """Yield (real_path, s3_key) for every file in the NIM cache."""
    cache = Path(cache_dir).expanduser() / "ngc/hub"

    if not cache.exists():
        raise SystemExit(f"NIM cache not found at {cache}")

    for snap in cache.glob("models--*/snapshots/*"):
        model = snap.parents[1].name[8:].replace("--", "/")
        sid = snap.name

        for f in filter(Path.is_file, snap.rglob("*")):
            real = f.resolve()
            # Pre-encode "/" in subdirectory paths so NIM can find them.
            # NIM double-encodes nested paths: "dir/file" → "dir%2Ffile" → "dir%252Ffile"
            file_path = str(f.relative_to(snap)).replace("/", "%2F")
            key = urllib.parse.quote(f"{model}:{sid}?file={file_path}", safe="")
            yield real, key


def s3_client(endpoint_url, access_key, secret_key, region, verify_ssl=True):
    import boto3
    from botocore.config import Config

    return boto3.client(
        "s3",
        endpoint_url=endpoint_url,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
        region_name=region,
        verify=verify_ssl,
        config=Config(
            signature_version="s3v4",
            s3={"addressing_style": "path"},
        ),
    )


def upload(cache_dir, bucket, endpoint_url, access_key, secret_key,
           region="us-east-1", verify_ssl=True, dry_run=False):
    pairs = list(iter_keys(cache_dir))
    total = len(pairs)
    print(f"{'[DRY RUN] ' if dry_run else ''}Uploading {total} files to bucket '{bucket}'")

    # Each (real_path, key) pair is a unique upload — even if two profiles share
    # the same blob (same real_path), they have different keys and both must exist in S3.
    # We only skip exact duplicates where both real_path AND key are identical.
    seen = set()
    s3 = None if dry_run else s3_client(
        endpoint_url, access_key, secret_key, region, verify_ssl
    )

    skipped = 0
    for i, (real, key) in enumerate(pairs, 1):
        if (real, key) in seen:
            print(f"  [{i}/{total}] SKIP (exact duplicate): {key}")
            skipped += 1
            continue

        seen.add((real, key))
        print(f"  [{i}/{total}] {key}")

        if not dry_run:
            s3.upload_file(str(real), bucket, key)

    uploaded = total - skipped
    print(f"\nDone. Uploaded: {uploaded}, Skipped (exact duplicates): {skipped}")


if __name__ == "__main__":
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    p = argparse.ArgumentParser()
    p.add_argument("--cache-dir", required=True, help="Path to the local NIM cache")
    p.add_argument("--bucket-name", required=True, help="Destination S3/MinIO bucket name")
    p.add_argument("--endpoint-url", default=None, help="S3/MinIO endpoint URL (optional, omit for AWS S3)")
    p.add_argument("--access-key", default=None, help="S3/MinIO access key (optional, falls back to env/IAM)")
    p.add_argument("--secret-key", default=None, help="S3/MinIO secret key (optional, falls back to env/IAM)")
    p.add_argument("--region", default="us-east-1", help="AWS region (default: us-east-1)")
    p.add_argument("--insecure", action="store_true", help="Skip SSL certificate verification")
    p.add_argument("--dry-run", action="store_true", help="Print keys without uploading anything")
    args = p.parse_args()

    upload(
        cache_dir=args.cache_dir,
        bucket=args.bucket_name,
        endpoint_url=args.endpoint_url,
        access_key=args.access_key,
        secret_key=args.secret_key,
        region=args.region,
        verify_ssl=not args.insecure,
        dry_run=args.dry_run,
    )