dexverse-cli 0.0.2__py3-none-any.whl

dexverse/downloader/manager.py

@@ -0,0 +1,365 @@
+"""Download manager — orchestrates skip/resume, backend dispatch, verify, state.
+
+Pipeline (doc §17):
+
+    for each file:
+        if dest exists and size matches  -> skip (resume)
+        else                             -> queue for backend
+    run backend on the queue
+    verify every file (size; sha256 under --deep)
+    write download-state.json + download.log
+
+State lives under ``<out>/.dexverse/tasks/<taskId>/`` and never touches the
+dataset tree.
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+from ..manifest import FileSpec, Manifest, resolve_dest
+from . import state as st
+from .base import Backend, ProgressReporter
+from .http_backend import HttpBackend, _HTTP_BACKEND_NAMES
+from .python_s3_backend import PythonS3Backend
+from .s5cmd_backend import S5CmdBackend, S5CmdNotAvailable
+from .verify import verify_files
+
+
+def _persist_manifest_copy(out: Path, task_id: str, manifest: Manifest) -> None:
+    """Write the full manifest copy (used by verify/repair) with mode 0600.
+
+    The copy may contain S3 credentials and/or presigned URLs, so it is
+    restricted to the owning user on POSIX (best-effort; no-op on Windows).
+    """
+    path = st.manifest_copy_path(out, task_id)
+    path.write_text(manifest.model_dump_json(indent=2), encoding="utf-8")
+    try:
+        os.chmod(path, 0o600)
+    except OSError:
+        pass
+
+
+def _sync_task_instructions(
+    out: Path,
+    manifest: Manifest,
+    task_state: st.TaskState,
+) -> None:
+    """Synchronize per-file task instructions after download/verification."""
+    successful_statuses = {st.DOWNLOADED, st.SKIPPED, st.VERIFIED}
+    for file_spec in manifest.files:
+        key = file_spec.effective_key()
+        file_state = task_state.files.get(key)
+        instruction_path = st.task_instruction_path(out, file_spec.relpath())
+        if (
+            file_state is not None
+            and file_state.status in successful_statuses
+            and file_spec.taskInstruction is not None
+        ):
+            instruction_path.parent.mkdir(parents=True, exist_ok=True)
+            instruction_path.write_text(
+                json.dumps(
+                    file_spec.taskInstruction,
+                    ensure_ascii=False,
+                    indent=2,
+                )
+                + "\n",
+                encoding="utf-8",
+            )
+        else:
+            instruction_path.unlink(missing_ok=True)
+
+
+def get_backend(name: str, **kwargs) -> Backend:
+    name = (name or "http").lower()
+    if name == "s5cmd":
+        return S5CmdBackend(**kwargs)
+    if name in ("python", "boto3"):
+        return PythonS3Backend(**kwargs)
+    if name in _HTTP_BACKEND_NAMES:
+        return HttpBackend(**kwargs)
+    raise ValueError(
+        f"unknown backend: {name!r} (expected 's5cmd', 'python', or 'http')"
+    )
+
+
+@dataclass
+class DownloadReport:
+    task_id: str
+    out: Path
+    backend: str
+    total: int
+    skipped: int
+    downloaded: int
+    failed: int
+    verified: int
+    corrupt: int
+    task_dir: Path
+    state_path: Path
+    failures: List[Dict[str, object]] = field(default_factory=list)
+
+    @property
+    def ok(self) -> bool:
+        return self.failed == 0 and self.corrupt == 0
+
+    def summary_line(self) -> str:
+        return (
+            f"task={self.task_id} backend={self.backend} "
+            f"total={self.total} downloaded={self.downloaded} "
+            f"skipped={self.skipped} failed={self.failed} "
+            f"verified={self.verified} corrupt={self.corrupt}"
+        )
+
+
+class DownloadManager:
+    def __init__(self, out: Union[str, Path], backend: str = "http",
+                 parallel: Optional[int] = None, *, verify: bool = True,
+                 deep: bool = False, persist_manifest: bool = True,
+                 insecure: bool = True,
+                 backend_kwargs: Optional[dict] = None,
+                 progress: Optional[ProgressReporter] = None):
+        self.out = Path(out).expanduser().resolve()
+        self.backend_name = backend
+        self.parallel = parallel
+        self.verify = verify
+        self.deep = deep
+        self.persist_manifest = persist_manifest
+        self.insecure = insecure
+        self.backend_kwargs = backend_kwargs or {}
+        self.progress = progress
+
+    # ------------------------------------------------------------------ #
+    def _resolve_parallel(self, manifest: Manifest) -> int:
+        if self.parallel is not None:
+            return self.parallel
+        return manifest.download.defaultParallel
+
+    def _ensure_storage(self, manifest: Manifest) -> None:
+        backend = (self.backend_name or "").lower()
+        if backend in _HTTP_BACKEND_NAMES:
+            # http backend uses per-file presigned URLs; no endpoint/creds needed.
+            missing = [f.effective_key() for f in manifest.files if not f.presignedUrl]
+            if missing:
+                raise ValueError(
+                    f"--backend http requires a presignedUrl per file; "
+                    f"{len(missing)} file(s) missing one "
+                    f"(first: {missing[0]!r})"
+                )
+            return
+        if manifest.storage is None or not manifest.storage.endpoint:
+            raise ValueError(
+                "manifest has no storage.endpoint — cannot download. "
+                "Provide a manifest with storage.endpoint/bucket, or set "
+                "AWS env vars, or use --backend http with per-file presignedUrl."
+            )
+
+    # ------------------------------------------------------------------ #
+    def run(self, manifest: Manifest) -> DownloadReport:
+        self._ensure_storage(manifest)
+        self.out.mkdir(parents=True, exist_ok=True)
+
+        # The on-disk manifest copy is the FULL task manifest (used by verify
+        # per §10.4). repair() reuses run() on a *subset* manifest, so it sets
+        # persist_manifest=False and writes the full manifest itself afterwards.
+        task_id = manifest.effective_task_id()
+        st.ensure_task_dir(self.out, task_id)
+        if self.persist_manifest:
+            _persist_manifest_copy(self.out, task_id, manifest)
+        log_path = st.log_file_path(self.out, task_id)
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+
+        backend_kwargs = dict(self.backend_kwargs)
+        if (self.backend_name or "").lower() in _HTTP_BACKEND_NAMES:
+            backend_kwargs.setdefault("verify", not self.insecure)
+        backend = get_backend(self.backend_name, **backend_kwargs)
+        parallel = self._resolve_parallel(manifest)
+
+        task_state = st.TaskState.load_or_create(self.out, manifest)
+        task_state.start(backend.name, parallel)
+        task_state.save()
+
+        # ---- decide what to fetch (resume / skip) ---------------------- #
+        todo: List[FileSpec] = []
+        n_skipped = 0
+        for f in manifest.files:
+            key = f.effective_key()
+            dest = resolve_dest(self.out, f.relpath())
+            file_state = task_state.files.get(key)
+            expected_size = f.size
+            if expected_size is None and file_state is not None:
+                expected_size = file_state.expectedSize
+            if (
+                dest.exists()
+                and (
+                    expected_size is None
+                    or dest.stat().st_size == expected_size
+                )
+            ):
+                n_skipped += 1
+                task_state.mark(key, st.SKIPPED,
+                                localSize=dest.stat().st_size,
+                                expectedSize=expected_size)
+            else:
+                todo.append(f)
+                task_state.mark(key, st.PENDING)
+        task_state.save()
+
+        # ---- run backend on the pending subset ------------------------- #
+        n_downloaded = n_failed = 0
+        if todo:
+            known_sizes = [f.size for f in todo if f.size is not None]
+            total_bytes = sum(known_sizes) if len(known_sizes) == len(todo) else None
+            if self.progress is not None:
+                self.progress.start(total_bytes=total_bytes, total_files=len(todo))
+            try:
+                results = backend.download(
+                    manifest, todo, self.out,
+                    parallel=parallel, log_path=log_path,
+                    progress=self.progress,
+                )
+            finally:
+                if self.progress is not None:
+                    self.progress.finish()
+            for f in todo:
+                key = f.effective_key()
+                r = results.get(key)
+                if r and r.ok:
+                    n_downloaded += 1
+                    dest = resolve_dest(self.out, f.relpath())
+                    task_state.mark(
+                        key, st.DOWNLOADED,
+                        localSize=dest.stat().st_size if dest.exists() else None,
+                        expectedSize=r.expectedSize,
+                    )
+                else:
+                    n_failed += 1
+                    err = r.error if r else "no result from backend"
+                    task_state.mark(
+                        key, st.FAILED,
+                        localSize=r.localSize if r else None,
+                        expectedSize=r.expectedSize if r else None,
+                        error=err,
+                    )
+            task_state.save()
+
+        # ---- verify ---------------------------------------------------- #
+        # Counts *actions taken this run*, not final statuses — verify promotes
+        # DOWNLOADED→VERIFIED, so reading task_state.counts() would hide downloads.
+        verified = corrupt = 0
+        if self.verify:
+            from .verify import sha256_of
+            for f in manifest.files:
+                key = f.effective_key()
+                # A file that failed to download stays FAILED — it is not
+                # "corrupt", and re-marking it would hide the real failure.
+                fs = task_state.files.get(key)
+                if fs and fs.status == st.FAILED:
+                    continue
+                dest = resolve_dest(self.out, f.relpath())
+                expected_size = f.size
+                if expected_size is None and fs is not None:
+                    expected_size = fs.expectedSize
+                if (
+                    not dest.exists()
+                    or (
+                        expected_size is not None
+                        and dest.stat().st_size != expected_size
+                    )
+                ):
+                    corrupt += 1
+                    task_state.mark(key, st.CORRUPT)
+                    continue
+                if self.deep and f.sha256 and \
+                        sha256_of(dest).lower() != f.sha256.lower():
+                    corrupt += 1
+                    task_state.mark(key, st.CORRUPT)
+                    continue
+                verified += 1
+                task_state.mark(key, st.VERIFIED)
+
+        _sync_task_instructions(self.out, manifest, task_state)
+        task_state.finish()
+        task_state.save()
+
+        failures = [
+            {"objectKey": fs.objectKey, "relpath": fs.relpath, "error": fs.error}
+            for fs in task_state.files.values()
+            if fs.status in (st.FAILED, st.CORRUPT)
+        ]
+
+        return DownloadReport(
+            task_id=task_id,
+            out=self.out,
+            backend=backend.name,
+            total=len(manifest.files),
+            skipped=n_skipped,
+            downloaded=n_downloaded,
+            failed=n_failed,
+            verified=verified,
+            corrupt=corrupt,
+            task_dir=st.task_dir(self.out, task_id),
+            state_path=st.state_file_path(self.out, task_id),
+            failures=failures,
+        )
+
+
+def repair(manifest: Manifest, out: Union[str, Path], *,
+           backend: str = "http", parallel: Optional[int] = None,
+           deep: bool = False, insecure: bool = True,
+           backend_kwargs: Optional[dict] = None,
+           progress: Optional[ProgressReporter] = None) -> DownloadReport:
+    """Re-fetch only the files that are missing or failed size/sha (doc §10.5).
+
+    Runs the download on the *subset* of broken files, but keeps the on-disk
+    task manifest as the FULL manifest so a later `verify` still covers every
+    file (§10.4). The subset manifest is pinned to the full task id so state
+    lands in the original task dir even when taskId is absent.
+    """
+    out_path = Path(out).expanduser().resolve()
+    full_task_id = manifest.effective_task_id()
+    try:
+        previous_state = st.TaskState.load(out_path, full_task_id)
+    except (OSError, ValueError, TypeError, json.JSONDecodeError):
+        previous_state = None
+
+    # 1. find broken files
+    broken: List[FileSpec] = []
+    for f in manifest.files:
+        key = f.effective_key()
+        dest = resolve_dest(out_path, f.relpath())
+        if not dest.exists():
+            broken.append(f)
+            continue
+        expected_size = f.size
+        if expected_size is None and previous_state is not None:
+            file_state = previous_state.files.get(key)
+            if file_state is not None:
+                expected_size = file_state.expectedSize
+        if expected_size is not None and dest.stat().st_size != expected_size:
+            broken.append(f)
+            continue
+        if deep and f.sha256:
+            from .verify import sha256_of
+            if sha256_of(dest).lower() != f.sha256.lower():
+                broken.append(f)
+
+    # 2. download the subset (persist_manifest=False — full manifest written below).
+    sub = manifest.model_copy(update={"files": broken})
+    if not sub.taskId:
+        sub.taskId = full_task_id  # avoid hash-fragmentation on the subset
+    mgr = DownloadManager(
+        out_path, backend=backend, parallel=parallel,
+        verify=True, deep=deep, persist_manifest=False, insecure=insecure,
+        backend_kwargs=backend_kwargs, progress=progress,
+    )
+    report = mgr.run(sub)
+    if not broken:
+        report.downloaded = 0
+
+    # 3. the on-disk manifest must remain the FULL manifest (verify covers all files)
+    st.ensure_task_dir(out_path, full_task_id)
+    _persist_manifest_copy(out_path, full_task_id, manifest)
+    return report

dexverse/downloader/python_s3_backend.py

@@ -0,0 +1,137 @@
+"""Pure-Python S3 backend (doc §4.2): boto3 fallback when s5cmd is unavailable.
+
+Downloads each object with a streaming ``get_object`` into ``out / relpath``.
+Parallelism is a thread pool (S3 I/O releases the GIL while waiting on the
+socket). boto3 is imported lazily so the package works without the ``[s3]``
+extra until this backend is actually selected.
+"""
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from ..manifest import FileSpec, Manifest, resolve_dest
+from .base import Credentials, FileResult, ProgressReporter
+
+
+# 8 MiB streaming chunks.
+_STREAM_CHUNK = 8 * 1024 * 1024
+
+
+class PythonS3Backend:
+    name = "python"
+
+    def __init__(self, *, client_factory=None, log_path: Optional[Path] = None):
+        # client_factory is for tests (e.g. moto). Signature: factory(creds) -> client
+        self._client_factory = client_factory
+        self._log_path = log_path
+
+    def _make_client(self, creds: Credentials):
+        if self._client_factory:
+            return self._client_factory(creds)
+
+        import boto3  # lazy import — needs the [s3] extra
+        from botocore.config import Config
+
+        cfg = Config(
+            s3={"addressing_style": "path"},
+            retries={"max_attempts": 5, "mode": "adaptive"},
+        )
+        kwargs = {
+            "service_name": "s3",
+            "endpoint_url": creds.endpoint,
+            "region_name": creds.region or "us-east-1",
+            "config": cfg,
+        }
+        if creds.access_key and creds.secret_key:
+            kwargs["aws_access_key_id"] = creds.access_key
+            kwargs["aws_secret_access_key"] = creds.secret_key
+            if creds.session_token:
+                kwargs["aws_session_token"] = creds.session_token
+        return boto3.client(**kwargs)
+
+    def download(
+        self,
+        manifest: Manifest,
+        files: List[FileSpec],
+        out: Path,
+        *,
+        parallel: int,
+        log_path: Optional[Path] = None,
+        progress: Optional[ProgressReporter] = None,
+    ) -> Dict[str, FileResult]:
+        if not files:
+            return {}
+
+        creds = Credentials.from_manifest(manifest)
+        if not creds.endpoint:
+            raise ValueError("manifest.storage.endpoint is required for the python backend")
+
+        default_bucket = manifest.storage.bucket if manifest.storage else None
+        client = self._make_client(creds)
+
+        workers = max(1, min(parallel or 4, 32))
+        results: Dict[str, FileResult] = {}
+
+        def _one(f: FileSpec) -> FileResult:
+            bucket = f.resolved_bucket(default_bucket)
+            if not bucket:
+                return FileResult(f.objectKey, False, error="no bucket for file")
+            dest = resolve_dest(out, f.relpath())
+            # Atomic landing: write to a sibling .part, replace dest only on
+            # success, so a failed get_object never leaves a truncated file.
+            tmp = dest.with_name(dest.name + ".part")
+            try:
+                tmp.parent.mkdir(parents=True, exist_ok=True)
+                obj = client.get_object(Bucket=bucket, Key=f.objectKey)
+                size = 0
+                with tmp.open("wb") as fh:
+                    while True:
+                        chunk = obj["Body"].read(_STREAM_CHUNK)
+                        if not chunk:
+                            break
+                        fh.write(chunk)
+                        size += len(chunk)
+                        if progress is not None:
+                            progress.advance(len(chunk))
+                if f.size is not None and size != f.size:
+                    tmp.unlink(missing_ok=True)
+                    return FileResult(
+                        f.objectKey, False,
+                        error=f"size mismatch: got {size}, expected {f.size}",
+                    )
+                tmp.replace(dest)
+                return FileResult(f.objectKey, True)
+            except Exception as exc:  # noqa: BLE001 — surface per-file failure
+                try:
+                    tmp.unlink(missing_ok=True)
+                except Exception:
+                    pass
+                return FileResult(f.objectKey, False, error=str(exc))
+
+        with ThreadPoolExecutor(max_workers=workers) as pool:
+            futures = [pool.submit(_one, f) for f in files]
+            for fut in as_completed(futures):
+                r = fut.result()
+                results[r.objectKey] = r
+                if progress is not None:
+                    progress.complete_file()
+
+        self._log(log_path, creds, files, results)
+        return results
+
+    def _log(self, log_path: Optional[Path], creds: Credentials,
+             files: List[FileSpec], results: Dict[str, FileResult]) -> None:
+        if not log_path:
+            return
+        log_path.parent.mkdir(parents=True, exist_ok=True)
+        ok = sum(1 for r in results.values() if r.ok)
+        with log_path.open("a", encoding="utf-8") as logf:
+            logf.write(
+                f"[python backend] endpoint={creds.endpoint} "
+                f"region={creds.region} files={len(files)} ok={ok}\n"
+            )
+            for r in results.values():
+                if not r.ok:
+                    logf.write(f"  FAIL {r.objectKey}: {r.error}\n")

dexverse/downloader/s5cmd_backend.py

@@ -0,0 +1,170 @@
+"""s5cmd backend (doc §7): generate ``download.s5cmd`` and run the binary.
+
+The run-file lists one ``cp`` per object, quoting both the s3 source and the
+local destination. Destinations are written with forward slashes (Go's s5cmd
+accepts them on Windows), which avoids fragile backslash escaping.
+
+Per-file success is determined *after* the run by checking each target exists
+with the expected size — this is robust regardless of how s5cmd formats its
+output, and lets repair target exactly the objects that failed.
+"""
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from ..manifest import FileSpec, Manifest, resolve_dest
+from .base import Credentials, FileResult, ProgressReporter
+
+
+class S5CmdNotAvailable(RuntimeError):
+    pass
+
+
+def find_s5cmd() -> Optional[str]:
+    """Locate the s5cmd binary: $DEXVERSE_S5CMD, then PATH, then common spots."""
+    override = os.getenv("DEXVERSE_S5CMD")
+    if override and Path(override).exists():
+        return override
+    found = shutil.which("s5cmd")
+    if found:
+        return found
+    candidates = [
+        "/usr/local/bin/s5cmd",
+        "/usr/bin/s5cmd",
+        str(Path.home() / ".local" / "bin" / "s5cmd"),
+    ]
+    for c in candidates:
+        if Path(c).exists():
+            return c
+    return None
+
+
+def _q(value: str) -> str:
+    """Quote an argument for the s5cmd run-file; escape backslash and quote."""
+    escaped = value.replace("\\", "\\\\").replace('"', '\\"')
+    return f'"{escaped}"'
+
+
+def generate_s5cmd_lines(manifest: Manifest, files: List[FileSpec], out: Path) -> List[str]:
+    """One ``cp`` line per file: cp "s3://bucket/key" "/out/rel/path"."""
+    default_bucket = manifest.storage.bucket if manifest.storage else None
+    lines: List[str] = []
+    for f in files:
+        bucket = f.resolved_bucket(default_bucket)
+        if not bucket:
+            raise ValueError(
+                f"file {f.objectKey!r} has no bucket (manifest.storage.bucket missing)"
+            )
+        src = f"s3://{bucket}/{f.objectKey.lstrip('/')}"
+        dst = resolve_dest(out, f.relpath()).as_posix()
+        lines.append(f"cp {_q(src)} {_q(dst)}")
+    return lines
+
+
+def write_s5cmd_file(manifest: Manifest, files: List[FileSpec], out: Path,
+                     dest_file: Path) -> Path:
+    lines = generate_s5cmd_lines(manifest, files, out)
+    dest_file.parent.mkdir(parents=True, exist_ok=True)
+    dest_file.write_text("\n".join(lines) + ("\n" if lines else ""), encoding="utf-8")
+    return dest_file
+
+
+class S5CmdBackend:
+    name = "s5cmd"
+
+    def __init__(self, s5cmd_path: Optional[str] = None):
+        self._s5cmd = s5cmd_path
+
+    def _resolve(self) -> str:
+        if self._s5cmd:
+            if Path(self._s5cmd).exists():
+                return self._s5cmd
+            raise S5CmdNotAvailable(
+                f"s5cmd not found at {self._s5cmd}. Set DEXVERSE_S5CMD to a valid "
+                "path, or use --backend python."
+            )
+        path = find_s5cmd()
+        if not path:
+            raise S5CmdNotAvailable(
+                "s5cmd binary not found. Install it "
+                "(https://github.com/peak/s5cmd), set DEXVERSE_S5CMD, "
+                "or use --backend python."
+            )
+        return path
+
+    def download(
+        self,
+        manifest: Manifest,
+        files: List[FileSpec],
+        out: Path,
+        *,
+        parallel: int,
+        log_path: Optional[Path] = None,
+        progress: Optional[ProgressReporter] = None,
+    ) -> Dict[str, FileResult]:
+        if not files:
+            return {}
+        s5cmd = self._resolve()
+
+        creds = Credentials.from_manifest(manifest)
+        if not creds.endpoint:
+            raise ValueError("manifest.storage.endpoint is required for the s5cmd backend")
+
+        # Write the run-file into the task dir beside the state file.
+        if log_path:
+            run_file = log_path.parent / "download.s5cmd"
+        else:
+            run_file = Path(out) / ".dexverse" / "_download.s5cmd"
+        write_s5cmd_file(manifest, files, out, run_file)
+
+        cmd = [
+            s5cmd,
+            "--endpoint-url", creds.endpoint,
+            "--numworkers", str(max(1, parallel)),
+            "run", str(run_file),
+        ]
+
+        env = dict(os.environ)
+        env.update(creds.as_env())
+
+        if log_path:
+            log_path.parent.mkdir(parents=True, exist_ok=True)
+            log_handle = log_path.open("a", encoding="utf-8")
+        else:
+            log_handle = open(os.devnull, "w", encoding="utf-8")
+        with log_handle as logf:
+            logf.write(f"$ {' '.join(cmd)}\n")
+            proc = subprocess.run(
+                cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                text=True,
+            )
+            logf.write(proc.stdout or "")
+            logf.write(f"\n[exit code {proc.returncode}]\n\n")
+
+        # Post-check per file: exists + size match ⇒ ok.
+        results = _postcheck(files, out, proc.returncode)
+        if progress is not None:
+            for f in files:
+                if results[f.objectKey].ok:
+                    dest = resolve_dest(out, f.relpath())
+                    progress.advance(dest.stat().st_size)
+                progress.complete_file()
+        return results
+
+
+def _postcheck(files: List[FileSpec], out: Path, exit_code: int) -> Dict[str, FileResult]:
+    results: Dict[str, FileResult] = {}
+    for f in files:
+        dest = resolve_dest(out, f.relpath())
+        if dest.exists() and (f.size is None or dest.stat().st_size == f.size):
+            results[f.objectKey] = FileResult(f.objectKey, True)
+        else:
+            results[f.objectKey] = FileResult(
+                f.objectKey, False,
+                error=f"s5cmd exit {exit_code}; {dest} missing or size mismatch",
+            )
+    return results