#!/usr/bin/env python3 # Copyright Advanced Micro Devices, Inc. # SPDX-License-Identifier: MIT """Fetch large binary artifacts referenced by .dvc pointer files from an S3 remote. Minimal in-tree replacement for `dvc pull` covering TheRock's only use of dvc: pulling MD5-content-addressed binary artifacts (MIOpen kernel databases, PAL prebuilt libs) from anonymous-public S3 buckets. Does not implement add/push/repro. Pointer file schema (the only one we handle): outs: - md5: <32-char-hex> # may end with `.dir` for directory hashes size: hash: md5 remote: # optional: per-output remote override (see below) path: Remote config (.dvc/config, INI): [core] remote = storage ['remote "storage"'] url = s3:/// allow_anonymous_login = true ['remote "golden-data"'] url = s3:/// allow_anonymous_login = true An `outs:` entry may carry an optional `remote: ` key that routes that specific output to a non-default remote (dvc's per-output remote routing). When present, `` must match a `['remote ""']` section; when absent the [core] default remote is used. Remotes are resolved per output, so a single pointer file (and a single pull) may pull blobs from several buckets/prefixes, each with its own `allow_anonymous_login` setting. S3 key scheme (dvc 3.x "new" layout, used by therock-dvc): `/files/md5//` for both files and `.dir` manifests. Library: from fetch_dvc_artifacts import pull result = pull(Path("/path/to/project")) CLI: python fetch_dvc_artifacts.py pull [DIR ...] [--jobs N] [--cache-dir PATH] [--no-cache] [-v] """ import argparse import concurrent.futures import configparser import hashlib import json import logging import os import shutil import sys import tempfile import uuid from dataclasses import dataclass from pathlib import Path from typing import Callable from urllib.parse import urlparse import boto3 from botocore import UNSIGNED from botocore.client import Config DEFAULT_JOBS = 8 DEFAULT_CACHE_DIR = Path.home() / ".cache" / "therock-dvc" class FetchError(Exception): """Raised on missing config, download failure, or MD5 mismatch.""" @dataclass class PullResult: fetched: int = 0 # downloaded from remote cached: int = 0 # served from local content-addressed cache skipped: int = 0 # already at destination with matching md5 def __iadd__(self, other: "PullResult") -> "PullResult": self.fetched += other.fetched self.cached += other.cached self.skipped += other.skipped return self @dataclass(frozen=True) class _Out: """One entry from a .dvc pointer file's `outs:` list. `remote` is the optional per-output remote name (dvc's per-output routing); None means "use the [core] default remote". """ md5: str size: int path: str remote: str | None = None @property def is_dir(self) -> bool: return self.md5.endswith(".dir") @property def bare_md5(self) -> str: return self.md5[:-4] if self.is_dir else self.md5 @dataclass(frozen=True) class _Remote: name: str bucket: str prefix: str anonymous: bool # -------------------------------------------------------------------------- # .dvc YAML subset parser # -------------------------------------------------------------------------- def _parse_dvc_pointer(path: Path) -> list[_Out]: """Parse a .dvc pointer file. Schema is fixed; full YAML is overkill.""" text = path.read_text() outs: list[_Out] = [] current: dict[str, str] | None = None in_outs = False for raw_line in text.splitlines(): line = raw_line.rstrip() if not line or line.lstrip().startswith("#"): continue if not in_outs: if line == "outs:": in_outs = True continue if line.startswith("- "): if current is not None: outs.append(_to_out(current, path)) current = {} line = line[2:] elif line.startswith(" "): line = line[2:] else: # A non-indented line ends the outs: section (e.g., another top-level key). in_outs = False continue if current is None: raise FetchError(f"{path}: 'outs:' content before list item: {raw_line!r}") if ":" not in line: raise FetchError(f"{path}: expected 'key: value', got {raw_line!r}") key, _, value = line.partition(":") current[key.strip()] = value.strip() if current is not None: outs.append(_to_out(current, path)) if not outs: raise FetchError(f"{path}: no 'outs:' entries") return outs def _to_out(d: dict[str, str], path: Path) -> _Out: try: md5 = d["md5"] size = int(d["size"]) out_path = d["path"] except KeyError as e: raise FetchError(f"{path}: missing required field {e}") from None except ValueError as e: raise FetchError(f"{path}: bad size: {e}") from None hash_alg = d.get("hash", "md5") if hash_alg != "md5": raise FetchError(f"{path}: unsupported hash algorithm {hash_alg!r}") bare = md5[:-4] if md5.endswith(".dir") else md5 if len(bare) != 32 or not all(c in "0123456789abcdef" for c in bare): raise FetchError(f"{path}: invalid md5 {md5!r}") remote = d.get("remote") or None return _Out(md5=md5, size=size, path=out_path, remote=remote) # -------------------------------------------------------------------------- # .dvc/config parser # -------------------------------------------------------------------------- @dataclass(frozen=True) class _RemoteConfig: """All remotes declared in .dvc/config, plus the [core] default name.""" default: str remotes: dict[str, _Remote] def resolve(self, name: str | None, pointer: Path) -> _Remote: """Return the _Remote for a per-output name (None -> default remote).""" key = name or self.default remote = self.remotes.get(key) if remote is None: raise FetchError( f"{pointer}: references remote {key!r} which is not defined in " f".dvc/config (known: {sorted(self.remotes)})" ) return remote def _canonical_remote_name(section: str) -> str | None: """Return the remote name if `section` is a `remote ""` header, else None. DVC's .dvc/config uses INI section headers like ['remote "storage"'] - the single quotes are literally part of the section name in the file, and Python's configparser preserves them. We strip the surrounding single quotes before matching, then extract the quoted name. """ s = section.strip("'") if not (s.startswith('remote "') and s.endswith('"')): return None return s[len('remote "') : -1] def _parse_dvc_remotes(path: Path) -> _RemoteConfig: """Parse every remote in .dvc/config and the [core] default remote name. Every declared remote is parsed eagerly (URL validated, anonymous flag read) so that a per-output `remote:` key can be resolved to any of them, and so a malformed remote surfaces up front rather than mid-download. """ cp = configparser.ConfigParser() cp.read(path) if "core" not in cp: raise FetchError(f"{path}: missing [core] section") default_name = cp["core"].get("remote") if not default_name: raise FetchError(f"{path}: [core] missing 'remote' key") remotes: dict[str, _Remote] = {} for section in cp.sections(): name = _canonical_remote_name(section) if name is None: continue url = cp[section].get("url") if not url: raise FetchError(f"{path}: [remote \"{name}\"] missing 'url'") parsed = urlparse(url) if parsed.scheme != "s3": raise FetchError(f"{path}: only s3:// remotes are supported, got {url!r}") anonymous = cp[section].getboolean("allow_anonymous_login", fallback=False) remotes[name] = _Remote( name=name, bucket=parsed.netloc, prefix=parsed.path.lstrip("/"), anonymous=anonymous, ) if default_name not in remotes: raise FetchError(f'{path}: missing [remote "{default_name}"] section') return _RemoteConfig(default=default_name, remotes=remotes) def _parse_dvc_config(path: Path) -> _Remote: """Return the [core] default remote. Thin wrapper over _parse_dvc_remotes.""" cfg = _parse_dvc_remotes(path) return cfg.remotes[cfg.default] # -------------------------------------------------------------------------- # S3 key + cache helpers # -------------------------------------------------------------------------- def _s3_key(remote: _Remote, md5: str) -> str: """Build the S3 key for a content-addressed blob. Uses the dvc 3.x "new" layout: /files/md5//. The therock-dvc bucket was migrated to this layout; the older // layout is not in use here. """ bare = md5[:-4] if md5.endswith(".dir") else md5 if len(bare) != 32: raise FetchError(f"invalid md5 length: {md5!r}") suffix = ".dir" if md5.endswith(".dir") else "" parts = [remote.prefix, "files", "md5", bare[:2], bare[2:] + suffix] return "/".join(p for p in parts if p) def _cache_path(cache_dir: Path, md5: str) -> Path: bare = md5[:-4] if md5.endswith(".dir") else md5 suffix = ".dir" if md5.endswith(".dir") else "" return cache_dir / bare[:2] / (bare[2:] + suffix) def _md5_of(path: Path) -> str: h = hashlib.md5() with path.open("rb") as f: for chunk in iter(lambda: f.read(1 << 20), b""): h.update(chunk) return h.hexdigest() def _materialize_from_cache(cache_file: Path, dest: Path) -> None: """Hardlink cache_file to dest; fall back to copy across filesystems.""" dest.parent.mkdir(parents=True, exist_ok=True) if dest.exists(): dest.unlink() try: os.link(cache_file, dest) except OSError: shutil.copy2(cache_file, dest) def _store_in_cache(src: Path, cache_file: Path) -> None: cache_file.parent.mkdir(parents=True, exist_ok=True) if cache_file.exists(): return # Unique temp name: distinct pointer files referencing the same content hash # are materialized by separate workers, which would otherwise race on a # shared `.tmp` path. os.replace stays atomic and idempotent. tmp = cache_file.with_suffix(f"{cache_file.suffix}.{uuid.uuid4().hex}.tmp") try: try: os.link(src, tmp) except OSError: shutil.copy2(src, tmp) try: os.replace(tmp, cache_file) except PermissionError: # Windows raises PermissionError when two workers race to os.replace # the same cache_file simultaneously (POSIX permits this silently). # If the winner already wrote the correct content, treat as success. if cache_file.exists() and src.stat().st_size == cache_file.stat().st_size: if _md5_of(src) == _md5_of(cache_file): return raise finally: if tmp.exists(): try: tmp.unlink() except OSError: pass # -------------------------------------------------------------------------- # Atomic-write download with MD5 verification # -------------------------------------------------------------------------- def _download_blob( s3, remote: _Remote, md5: str, dest: Path, expected_size: int | None ) -> None: """Download s3:/// to dest atomically, verifying md5 (and size). MD5 is computed by re-reading the file after download, not streamed during write. boto3's multipart download writes chunks out of order via seek+write, which breaks any sequential-hash wrapper. Re-reading costs one extra pass (~150 MB/s on SSD); negligible vs network time for our workload. """ key = _s3_key(remote, md5) dest.parent.mkdir(parents=True, exist_ok=True) tmp = dest.with_suffix(dest.suffix + ".tmp") if tmp.exists(): tmp.unlink() try: with tmp.open("wb") as f: s3.download_fileobj(remote.bucket, key, f) actual_size = tmp.stat().st_size if expected_size is not None and actual_size != expected_size: raise FetchError( f"{dest}: size mismatch (expected {expected_size}, got {actual_size})" ) actual_md5 = _md5_of(tmp) bare_expected = md5[:-4] if md5.endswith(".dir") else md5 if actual_md5 != bare_expected: raise FetchError( f"{dest}: md5 mismatch (expected {bare_expected}, got {actual_md5})" ) os.replace(tmp, dest) finally: if tmp.exists(): try: tmp.unlink() except OSError: pass # -------------------------------------------------------------------------- # Per-pointer-entry materialization # -------------------------------------------------------------------------- def _human(n: float) -> str: for unit in ("B", "KiB", "MiB", "GiB"): if n < 1024: return f"{n:.1f} {unit}" n /= 1024 return f"{n:.1f} TiB" def _materialize_file( s3, remote: _Remote, md5: str, size: int | None, dest: Path, cache_dir: Path | None, log: Callable[[str], None], ) -> PullResult: """Materialize one file: destination check -> cache check -> download. `size` may be None for files referenced from a .dir manifest (which carries md5 only). In that case the size check is skipped; md5 verification still happens. """ label = f"{dest.name}" if size is None else f"{dest.name} ({_human(size)})" # Fast path 1: destination already correct. if dest.exists(): if size is None or dest.stat().st_size == size: if _md5_of(dest) == md5: log(f" ok {label} [present]") return PullResult(skipped=1) # Fast path 2: content-addressed cache hit. if cache_dir is not None: cache_file = _cache_path(cache_dir, md5) if cache_file.exists(): if size is None or cache_file.stat().st_size == size: if _md5_of(cache_file) == md5: _materialize_from_cache(cache_file, dest) log(f" cached {label}") return PullResult(cached=1) log(f" warn cache entry corrupt for {md5[:8]}..., discarding") cache_file.unlink() log(f" fetching {label}") _download_blob(s3, remote, md5, dest, expected_size=size) if cache_dir is not None: _store_in_cache(dest, _cache_path(cache_dir, md5)) return PullResult(fetched=1) def _materialize_dir( s3, remote: _Remote, out: _Out, dest: Path, cache_dir: Path | None, log: Callable[[str], None], ) -> PullResult: """Handle a `.dir` directory hash: fetch JSON manifest, materialize each entry.""" log(f" manifest {dest}/ ({out.md5[:8]}...)") # Try cache for the manifest first. manifest_bytes: bytes | None = None cache_file = _cache_path(cache_dir, out.md5) if cache_dir is not None else None if cache_file is not None and cache_file.exists(): if _md5_of(cache_file) == out.bare_md5: manifest_bytes = cache_file.read_bytes() if manifest_bytes is None: # Download manifest. If we have a cache, write it there directly (it's # an immutable content-addressed blob). Otherwise, use a tempfile. if cache_file is not None: _download_blob(s3, remote, out.md5, cache_file, expected_size=None) manifest_bytes = cache_file.read_bytes() else: with tempfile.TemporaryDirectory(prefix="therock-dvc-") as tmp_dir: manifest_path = Path(tmp_dir) / "manifest" _download_blob(s3, remote, out.md5, manifest_path, expected_size=None) manifest_bytes = manifest_path.read_bytes() try: entries = json.loads(manifest_bytes) except json.JSONDecodeError as e: raise FetchError(f"{dest}: manifest is not valid JSON: {e}") from None if not isinstance(entries, list): raise FetchError(f"{dest}: manifest is not a JSON array") result = PullResult() for entry in entries: if not isinstance(entry, dict) or "md5" not in entry or "relpath" not in entry: raise FetchError(f"{dest}: malformed manifest entry {entry!r}") # Manifest entries don't carry size; trust md5 verification alone. result += _materialize_file( s3, remote, md5=entry["md5"], size=None, dest=dest / entry["relpath"], cache_dir=cache_dir, log=log, ) return result # -------------------------------------------------------------------------- # Discovery + parallel orchestration # -------------------------------------------------------------------------- def _walk_dvc_pointers(project_dir: Path) -> list[Path]: # rglob("*.dvc") matches the .dvc/ config directory too; filter to files. return sorted(p for p in project_dir.rglob("*.dvc") if p.is_file()) # Inner s3transfer max_concurrency is 10 by default; we bump the boto3 # connection pool to `jobs * INNER_S3TRANSFER_CONCURRENCY` so that multipart # ranged GETs don't starve when many outer workers hit S3 simultaneously. _INNER_S3TRANSFER_CONCURRENCY = 10 _MIN_POOL_CONNECTIONS = 50 _BOTO3_RETRIES = {"max_attempts": 5, "mode": "adaptive"} def _make_s3_client(remote: _Remote, *, max_pool_connections: int): """Build an S3 client. Anonymous (UNSIGNED) when the remote allows it.""" config_kwargs: dict[str, object] = { "max_pool_connections": max_pool_connections, "retries": _BOTO3_RETRIES, } if remote.anonymous: config_kwargs["signature_version"] = UNSIGNED return boto3.client("s3", config=Config(**config_kwargs)) def _process_pointer( cfg: _RemoteConfig, client_for: Callable[[_Remote], object], pointer: Path, cache_dir: Path | None, log: Callable[[str], None], ) -> PullResult: """Materialize every entry in one .dvc pointer file. Each entry's remote is resolved independently from its `remote:` key (falling back to the [core] default), so one pointer file may pull from several remotes. Children of a `.dir` manifest inherit the parent output's remote. """ outs = _parse_dvc_pointer(pointer) result = PullResult() for out in outs: remote = cfg.resolve(out.remote, pointer) s3 = client_for(remote) dest = pointer.parent / out.path if out.is_dir: result += _materialize_dir(s3, remote, out, dest, cache_dir, log) else: result += _materialize_file( s3, remote, md5=out.md5, size=out.size, dest=dest, cache_dir=cache_dir, log=log, ) return result def pull( project_dir: Path, *, jobs: int = DEFAULT_JOBS, cache_dir: Path | None = DEFAULT_CACHE_DIR, log: Callable[[str], None] = print, ) -> PullResult: """Materialize every .dvc-tracked file under project_dir from its S3 remote. Raises FetchError if .dvc/config is missing or any download fails. """ project_dir = Path(project_dir).resolve() config_file = project_dir / ".dvc" / "config" if not config_file.exists(): raise FetchError(f"no .dvc/config in {project_dir}") cfg = _parse_dvc_remotes(config_file) pointers = _walk_dvc_pointers(project_dir) if not pointers: log(f"no .dvc pointer files found under {project_dir}") return PullResult() # One S3 client per configured remote, built up front (before any worker # starts) so the dict is read-only during the parallel phase - no locking # needed. Clients are shared across workers (boto3 clients are thread-safe); # they can't be shared across remotes because the anonymous flag differs. pool_size = max(jobs * _INNER_S3TRANSFER_CONCURRENCY, _MIN_POOL_CONNECTIONS) clients = { name: _make_s3_client(remote, max_pool_connections=pool_size) for name, remote in cfg.remotes.items() } def client_for(remote: _Remote) -> object: return clients[remote.name] remotes_desc = ", ".join( f"s3://{r.bucket}/{r.prefix}" for r in cfg.remotes.values() ) log(f"pull: {len(pointers)} pointer file(s) from {remotes_desc}") result = PullResult() errors: list[tuple[Path, Exception]] = [] with concurrent.futures.ThreadPoolExecutor(max_workers=jobs) as ex: futures = { ex.submit(_process_pointer, cfg, client_for, p, cache_dir, log): p for p in pointers } for fut in concurrent.futures.as_completed(futures): pointer = futures[fut] try: result += fut.result() except Exception as e: errors.append((pointer, e)) if errors: for pointer, err in errors: log(f" ERROR {pointer}: {err}") raise FetchError(f"{len(errors)} file(s) failed to fetch (see log above)") return result # -------------------------------------------------------------------------- # CLI # -------------------------------------------------------------------------- def _main(argv: list[str]) -> int: parser = argparse.ArgumentParser( prog="fetch_dvc_artifacts", description="Fetch large binary artifacts referenced by .dvc files from S3.", ) sub = parser.add_subparsers(dest="cmd", required=True) p_pull = sub.add_parser("pull", help="Pull all .dvc-referenced artifacts.") p_pull.add_argument( "directories", nargs="*", default=["."], help="Project directories to pull (each must contain .dvc/config). " "Default: cwd.", ) p_pull.add_argument( "--jobs", type=int, default=DEFAULT_JOBS, help=f"Parallel pointer-file workers (default {DEFAULT_JOBS}).", ) p_pull.add_argument( "--cache-dir", type=Path, default=DEFAULT_CACHE_DIR, help=f"Content-addressed cache directory (default {DEFAULT_CACHE_DIR}).", ) p_pull.add_argument( "--no-cache", action="store_true", help="Disable the local content-addressed cache.", ) p_pull.add_argument( "-v", "--verbose", action="store_true", help="Enable boto3/botocore debug logging.", ) args = parser.parse_args(argv) if args.verbose: boto3.set_stream_logger("botocore", logging.INFO) cache_dir = None if args.no_cache else args.cache_dir total = PullResult() try: for d in args.directories: total += pull(Path(d), jobs=args.jobs, cache_dir=cache_dir) except FetchError as e: print(f"fetch_dvc_artifacts: {e}", file=sys.stderr) return 1 print( f"summary: fetched={total.fetched} " f"cached={total.cached} skipped={total.skipped}" ) return 0 if __name__ == "__main__": sys.exit(_main(sys.argv[1:]))