diff --git a/docs/commands.md b/docs/commands.md index a573327..095a5d7 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -206,9 +206,10 @@ Download DVC-tracked data without creating tracking files. ```bash dt get [-o ] [--rev REV] [--link TYPES] [-j N] [-f] [-v] dt get --csv [--path-col COL] [-o ] +dt get --csv -o s3://// [-w N] [--dest-profile NAME] ``` -The `dvc get` counterpart to `dt import`: materialises data and writes no `.dvc` file, for handing a subset of a dataset to someone outside the group. A path may name a subdirectory *inside* a tracked directory output, so you can take part of a large dataset without transferring all of it. With `--csv` the source is cloned once and every row resolved against that clone, rather than re-cloning per path. [Full documentation →](get.md) +The `dvc get` counterpart to `dt import`: materialises data and writes no `.dvc` file, for handing a subset of a dataset to someone outside the group. A path may name a subdirectory *inside* a tracked directory output, so you can take part of a large dataset without transferring all of it. With `--csv` the source is cloned once and every row resolved against that clone, rather than re-cloning per path. `-o s3://…` streams from the source's remote straight into object storage, and `-w N` spreads that stream across N compute nodes via qxub — network-bound work, so nodes add bandwidth. [Full documentation →](get.md) --- diff --git a/docs/config_options.md b/docs/config_options.md index 7cdb2f4..db8aabf 100644 --- a/docs/config_options.md +++ b/docs/config_options.md @@ -190,7 +190,7 @@ dt config set index.retry_interval 10 ## qxub Options -These options configure the jobs `dt` submits via [qxub](https://github.com/swarbricklab/qxub): `dt push --workers`, `dt add`, and (via the `archive.qxub_*` fallbacks) `dt remote archive stage --via-qxub`. +These options configure the jobs `dt` submits via [qxub](https://github.com/swarbricklab/qxub): `dt push --workers`, `dt get --workers`, `dt remote verify --workers`, `dt add`, and (via the `archive.qxub_*` fallbacks) `dt remote archive stage --via-qxub`. ### `qxub.env` @@ -284,7 +284,8 @@ dt push -w 16 ``` `dt pull` has no `--workers` flag; it pulls in-process. The `qxub.*` settings -apply to `dt push -w` and to `dt add` jobs submitted via qxub. +apply to `dt push -w`, `dt get -w` (`s3://` destinations only), `dt remote +verify -w`, and to `dt add` jobs submitted via qxub. ## auth Options diff --git a/docs/get.md b/docs/get.md index 07eb144..d52019b 100644 --- a/docs/get.md +++ b/docs/get.md @@ -113,6 +113,8 @@ S3 destinations only (see [To an S3 bucket](#to-an-s3-bucket)): - `--dest-region `: Destination region. Required if the profile's region is `auto`. - `--dest-account-id `: Abort unless the destination credentials resolve to this account - `--chunk-size `: Streaming chunk size (default: 8) +- `-w, --workers `: Distribute across n compute nodes via qxub — see [Across several compute nodes](#across-several-compute-nodes---workers) +- `--no-wait`: With `--workers`, submit the jobs and exit without waiting ## To an S3 bucket @@ -172,6 +174,71 @@ keep accruing storage charges. Add a lifecycle rule expiring incomplete multipar Not supported: `ssh://` destinations, and `--link` (there is nothing to link — an S3 destination is always a stream). +## Across several compute nodes: `--workers` + +`-w/--workers N` spreads an `s3://` transfer over N compute nodes via [qxub](https://github.com/swarbricklab/qxub), instead of N threads in one process. + +```bash +# 82 samples out to a collaborator's bucket, over 8 nodes +dt get bcarc_wts --csv samples.csv -o s3://their-bucket/fastqs/ -w 8 --dest-profile their-aws + +# Submit and go; collect the result later with qxub monitor +dt get bcarc_wts --csv samples.csv -o s3://their-bucket/fastqs/ -w 8 --no-wait +``` + +### Why only `s3://` + +`--workers` is an error on a local destination, and that is deliberate rather than unfinished: + +| | bound by | second node adds | +|---|---|---| +| `-o s3://…` | the network | **bandwidth** — a separate NIC | +| `-o some/dir/` | the filesystem both ends share | contention | + +Streaming into object storage is network-bound, so nodes multiply throughput. A local destination on Lustre is limited by the filesystem, not the node, so a second node competes with the first rather than helping. Raise `-j` there. + +On NCI there is a second reason: `copyq` is the queue with outbound network access, and it is already the default for `dt` transfer workers (`qxub.queue`). + +> **Worth measuring before you scale up.** If your *source* remote is on `/g/data` rather than object storage, the read side is a shared filesystem again and the same argument applies to it. Time a single-node `-j 8` run over a few samples and watch whether throughput saturates before adding nodes. + +### The two axes + +`--jobs` and `--workers` are not alternatives, and they compose — `-w 8 -j 4` is 8 nodes running 4 streams each. + +| | `-w, --workers` | `-j, --jobs` | +|---|---|---| +| Unit | compute nodes via qxub | threads in one process | +| Partitioned | up front, by size | not at all — a shared queue | + +The difference in the second row is the whole reason `--workers` needs more machinery than a bigger `-j`. Separate processes on separate machines cannot steal work from each other, so the split has to be decided before they start; a thread pool hands tasks out from a shared queue and self-corrects as it goes. + +### How the work is split + +Files are partitioned by **size**, largest-first onto the least-loaded worker (greedy LPT bin-packing — the same code `dt push --workers` uses). Fastq sample sizes are genuinely uneven, which is exactly the case this is for: a split by row count would hand one node a 40 GiB sample and another a 4 GiB one. + +The unit is a *file*, not a row, so one sample directory may be spread over several nodes. Reporting is unaffected — each file carries its row, and the per-row `✓`/`✗` report comes back in CSV order as usual: + +``` +Destination: s3://their-bucket/fastqs +Writing as: account 123456789012 as arn:aws:iam::123456789012:user/handoff +Partition byte balance across 8 active worker(s): min=41.2G, max=46.8G, total=358G +✓ data/fq/AF013-A: 2 files -> s3://their-bucket/fastqs/AF013-A +✓ data/fq/AF013-B: 2 files -> s3://their-bucket/fastqs/AF013-B +``` + +Read the balance line: `min` far below `max` means one file dominates a partition and more nodes will not help. + +### What each worker needs + +The submitting node resolves every row once — it has to, because the byte sizes the partitioning needs come from that resolve — and ships the resolved md5, size and destination key for each file in the manifest under `.dt/tmp/get//`. Consequences worth knowing: + +- **No worker runs `dvc list`.** DVC keeps a SQLite state database for the repository it reads, and pointing N compute nodes at the shared clone's copy of it is how you get `database is locked`. Each worker opens the clone once, read-only, for one thing only: the source remote's filesystem and credentials out of its `.dvc/config`. +- **The manifest directory and the clone must be visible from the compute nodes.** On `/scratch` or `/g/data` they are. This is the same requirement `dt push --workers` has. +- **Each worker preflights for itself**, so `--dest-account-id` is enforced on the node that actually writes — a credential chain can resolve differently on a compute node than on a login node. +- **No key material is written to the manifest.** The profile *name*, endpoint and region travel; secrets stay in `~/.aws`, which the compute nodes read for themselves. + +A worker killed mid-partition — walltime, usually — reports its rows as failed rather than silently short. Re-run with `--resume` to send only what is missing. + ## Examples ```bash diff --git a/dt/__init__.py b/dt/__init__.py index 67e64dd..ba35cfb 100644 --- a/dt/__init__.py +++ b/dt/__init__.py @@ -1,3 +1,3 @@ """DVC Tools - Convenient tools for working with DVC in HPC environments.""" -__version__ = "0.24.0" +__version__ = "0.25.0" diff --git a/dt/cli.py b/dt/cli.py index c1de370..f4887f5 100644 --- a/dt/cli.py +++ b/dt/cli.py @@ -838,6 +838,89 @@ def _report_rows(results, verb): raise SystemExit(1) +def _get_distributed(repository, path, out, csv_path, path_col, owner, rev, + num_workers, jobs, force, resume, check, refresh, + dest_config, chunk_bytes, to_s3, wait, verbose): + """Run ``dt get --workers``: submit qxub workers and report what landed. + + Split out of ``get_cmd`` because it is a third top-level mode alongside + single-path and ``--csv``, not a variation on either -- it consumes both + kinds of target list and then does something structurally different with + them. + """ + if num_workers < 1: + raise click.UsageError("--workers must be at least 1.") + + if not to_s3: + raise click.UsageError( + "--workers only applies to an s3:// destination.\n" + "Streaming to object storage is bound by the network, so more nodes " + "mean more bandwidth. A local destination is bound by the filesystem " + "both ends share, where a second node adds contention rather than " + "throughput -- raise -j instead." + ) + + # Rows land under the root prefix by name, matching the non-distributed + # path: a trailing slash on -o collects each row under its own basename, + # and a single path without one is used verbatim. + if csv_path: + if path: + raise click.UsageError("Do not provide PATH when using --csv.") + try: + rows = utils.read_csv_targets(csv_path, path_col, None) + except ValueError as e: + raise click.ClickException(str(e)) + if not rows: + raise click.ClickException(f"No rows selected from {csv_path}") + targets = [ + (row_path, row_out or Path(row_path or '.').name) + for row_path, row_out in rows + ] + else: + if not path: + raise click.UsageError("PATH is required (unless using --csv).") + targets = [(path, Path(path).name if out.endswith('/') else None)] + + # Preflight on the submitting node before spending a queue slot: a bad + # profile or a missing s3:PutObject grant should fail here, not eight times + # over in eight PBS jobs. + root = get_dest.S3Dest(out, dest_config) + identity = root.preflight() + click.echo(f"Destination: {root}") + click.echo(f"Writing as: {identity}") + + try: + job_ids, manifest_dir, results = get_mod.distributed_upload_to_s3( + repository=repository, + targets=targets, + root=root, + num_workers=num_workers, + owner=owner, + rev=rev, + jobs=jobs, + force=force, + resume=resume, + check=check, + chunk_size=chunk_bytes, + refresh=refresh, + path_col=path_col, + wait=wait, + verbose=verbose, + ) + except get_mod.GetError as e: + raise click.ClickException(str(e)) + + if results is None: + click.echo(f"Submitted {len(job_ids)} job(s):") + for job_id in job_ids: + click.echo(f" {job_id}") + click.echo(f"Manifest: {manifest_dir}") + click.echo(f"\nMonitor with: qxub monitor {' '.join(job_ids)}") + return + + _report_rows(results, "uploaded") + + def _bool_config(key, default): """Read a boolean config value, tolerating yes/no/true/false strings.""" raw = cfg.get_value(key) @@ -4574,7 +4657,7 @@ def import_cmd(repository, path, out, owner, no_checkout, no_refresh, no_downloa @cli.command('get') -@click.argument('repository') +@click.argument('repository', required=False) @click.argument('path', required=False) @click.option('-o', '--out', help='Destination path. With --csv, the directory rows are collected under') @click.option('--owner', help='Override the GitHub owner for short names') @@ -4593,11 +4676,19 @@ def import_cmd(repository, path, out, owner, no_checkout, no_refresh, no_downloa @click.option('--dest-region', default=None, help='Region for an s3:// destination. Required if the profile\'s region is "auto".') @click.option('--dest-account-id', default=None, help='Abort unless the destination credentials resolve to this AWS account. For unattended runs.') @click.option('--chunk-size', type=int, default=None, help='Streaming chunk size in MiB for s3:// destinations (default: 8)') +@click.option('--workers', '-w', type=int, default=None, + help='Distribute an s3:// transfer across N compute nodes via qxub.') +@click.option('--no-wait', is_flag=True, + help='With --workers: submit jobs and exit without waiting.') +@click.option('--worker', type=int, default=None, + help='Worker ID (internal, used by submitted jobs).') +@click.option('--manifest', type=click.Path(exists=True), default=None, + help='Manifest directory (internal, used by submitted jobs).') @click.option('-v', '--verbose', is_flag=True, help='Show detailed progress') def get_cmd(repository, path, out, owner, rev, csv_path, path_col, jobs, link, no_refresh, remote_fallback, resume, check, force, dest_profile, dest_endpoint_url, dest_region, dest_account_id, - chunk_size, verbose): + chunk_size, workers, no_wait, worker, manifest, verbose): """Download DVC-tracked data without creating tracking files. Wraps the idea of ``dvc get``: materialises data from a source repository @@ -4617,6 +4708,17 @@ def get_cmd(repository, path, out, owner, rev, csv_path, path_col, Unlike a loop over ``dvc get``, the source repository is cloned once and every row is resolved against that clone. + \b + Two kinds of parallelism, and they are not alternatives: + -j, --jobs threads inside one process + -w, --workers compute nodes via qxub; s3:// destinations only + + ``--workers`` exists because streaming into object storage is bound by the + network, so more nodes mean more bandwidth. A local destination is bound by + the filesystem both ends share, where a second node adds contention rather + than throughput -- widen ``-j`` there instead. Files are partitioned across + workers by size, so each gets a comparable share of the bytes. + \b Examples: dt get my-registry data/fq/AF013-A -o fastqs/ @@ -4624,7 +4726,30 @@ def get_cmd(repository, path, out, owner, rev, csv_path, path_col, dt get my-registry --csv samples.csv --path-col fq_dir -o fastqs/ dt get my-registry data/ref.fa --link copy -o ./ dt get my-registry data/fq/AF013-A -o s3://bucket/fastqs/ --dest-profile aws + dt get my-registry --csv samples.csv -o s3://bucket/fastqs/ -w 8 + dt get my-registry --csv samples.csv -o s3://bucket/fastqs/ -w 8 --no-wait """ + # Worker mode: everything this needs is in the manifest, including the + # destination and the clone, so no REPOSITORY argument is required or read. + if worker is not None: + if manifest is None: + raise click.UsageError("--worker requires --manifest.") + try: + written, failed = get_mod.worker_upload_to_s3( + manifest_dir=Path(manifest), + worker_id=worker, + verbose=verbose, + ) + except get_mod.GetError as e: + raise click.ClickException(str(e)) + click.echo(f"{written} uploaded, {failed} failed") + if failed: + raise SystemExit(1) + return + + if not repository: + raise click.UsageError("REPOSITORY is required.") + if force and resume: raise click.UsageError( "--force and --resume are contradictory: one re-fetches everything, " @@ -4662,6 +4787,20 @@ def get_cmd(repository, path, out, owner, rev, csv_path, path_col, chunk_bytes = (chunk_size * 1024 * 1024) if chunk_size \ else get_dest.DEFAULT_CHUNK_SIZE + if no_wait and workers is None: + raise click.UsageError("--no-wait only means anything with --workers.") + + if workers is not None: + _get_distributed( + repository=repository, path=path, out=out, csv_path=csv_path, + path_col=path_col, owner=owner, rev=rev, num_workers=workers, + jobs=jobs, force=force, resume=resume, check=check, + refresh=not no_refresh, dest_config=dest_config, + chunk_bytes=chunk_bytes, to_s3=to_s3, wait=not no_wait, + verbose=verbose, + ) + return + if csv_path: if path: raise click.UsageError("Do not provide PATH when using --csv.") diff --git a/dt/get.py b/dt/get.py index 3610580..0c5bb8a 100644 --- a/dt/get.py +++ b/dt/get.py @@ -18,16 +18,19 @@ of 412 samples without transferring the other 330. """ +import dataclasses import json import re import subprocess import time +import uuid from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Dict, List, Optional, Sequence, Tuple from . import cache_ops from . import get_dest +from . import hpc from . import remote as remote_mod from . import tmp as tmp_mod from . import utils @@ -207,6 +210,47 @@ def list_source_files( return files +def _resolve_rows( + clone: Path, + row_paths: Sequence[str], + rev: Optional[str], + jobs: int, + path_col: str = 'path', +) -> List[Tuple[Optional[List[Dict[str, object]]], Optional[str]]]: + """Resolve each row's path to a file listing, concurrently. + + Returns one ``(entries, error)`` pair per input path, in input order; + exactly one of the two is ever set. Destinations are deliberately *not* + computed here -- a local row lands in a directory and an S3 row lands under + a key prefix, and the callers differ on that while agreeing on this. + + Each row costs a ``dvc list`` subprocess, so they run concurrently: on a + long manifest the resolves are otherwise a serial prologue to all the real + work. + + ...but the first row is resolved on its own first. DVC keeps a SQLite state + db for the repository it is reading, and on a freshly-cloned one it does not + exist yet; fan out immediately and the workers race to create it, which + fails all but one with "database is locked". One serial call builds it, + after which concurrent readers are fine. + """ + def resolve(row_path: str): + if not row_path: + return None, f'Missing {path_col}' + try: + return list_source_files(clone, row_path, rev=rev), None + except GetError as e: + return None, str(e) + + resolved: List[Tuple[Optional[List[Dict[str, object]]], Optional[str]]] = [] + if row_paths: + resolved.append(resolve(row_paths[0])) + if len(row_paths) > 1: + with ThreadPoolExecutor(max_workers=max(1, jobs)) as pool: + resolved.extend(pool.map(resolve, row_paths[1:])) + return resolved + + def verify_file(dest: Path, md5: str) -> bool: """Does the file on disk hash to the value DVC recorded for it?""" try: @@ -534,6 +578,29 @@ def _upload_one( return relpath, True, 'uploaded' +def _upload_all( + tasks: Sequence[Tuple[Dict[str, object], 'get_dest.S3Dest']], + odb, + jobs: int, + force: bool, + resume: bool = False, + check: bool = False, + chunk_size: int = get_dest.DEFAULT_CHUNK_SIZE, +) -> List[Tuple[str, bool, str]]: + """Stream ``(entry, dest)`` pairs through one pool, preserving order. + + The S3 counterpart to :func:`_place_all`, and flat for the same reason: a + row is often a sample directory of two fastqs, so a pool per row would idle + most of the budget. + """ + def place(task): + entry, row_dest = task + return _upload_one(entry, odb, row_dest, force, resume, check, chunk_size) + + with ThreadPoolExecutor(max_workers=max(1, jobs)) as pool: + return list(pool.map(place, tasks)) + + def upload_to_s3( clone: Path, targets: Sequence[Tuple[str, 'get_dest.S3Dest']], @@ -555,27 +622,18 @@ def upload_to_s3( pool, for the same reason the local path does: a row is often a sample directory of two fastqs, and batching per row would idle most workers. - Unlike the local network path there is no ``dvc get`` subprocess and so no - contention on the clone's SQLite state database, which is what forced that - path to run rows serially. + Unlike the local network path there is no ``dvc get`` subprocess, so nothing + here forces rows to run serially. """ repo, odb = _open_source_odb(clone) try: - resolved = [] - for row_path, row_dest in targets: - if not row_path: - resolved.append((row_path, row_dest, None, 'missing source path')) - continue - try: - resolved.append( - (row_path, row_dest, list_source_files(clone, row_path, rev=rev), None) - ) - except GetError as e: - resolved.append((row_path, row_dest, None, str(e))) + resolved = _resolve_rows( + clone, [row_path for row_path, _ in targets], rev, jobs + ) tasks: List[Tuple[Dict[str, object], 'get_dest.S3Dest']] = [] spans: Dict[int, Tuple[int, int]] = {} - for i, (_, row_dest, entries, _) in enumerate(resolved): + for i, ((_, row_dest), (entries, _)) in enumerate(zip(targets, resolved)): if entries is None: continue start = len(tasks) @@ -587,17 +645,14 @@ def upload_to_s3( print(f"Streaming {len(tasks)} objects across {max(1, jobs)} workers " f"(~{mib} MiB peak memory)") - def place(task): - entry, row_dest = task - return _upload_one( - entry, odb, row_dest, force, resume, check, chunk_size - ) - - with ThreadPoolExecutor(max_workers=max(1, jobs)) as pool: - placed = list(pool.map(place, tasks)) + placed = _upload_all( + tasks, odb, jobs, force, resume, check, chunk_size + ) results: List[Tuple[str, bool, str]] = [] - for i, (row_path, row_dest, entries, error) in enumerate(resolved): + for i, ((row_path, row_dest), (entries, error)) in enumerate( + zip(targets, resolved) + ): if entries is None: results.append((row_path or '(empty)', False, error)) continue @@ -615,6 +670,347 @@ def place(task): repo.close() +# ============================================================================= +# Distributed streaming to S3 via qxub +# ============================================================================= +# +# Only the S3 destination has a distributed mode, and the reason is bandwidth +# rather than convenience. Streaming from a DVC remote into object storage is +# network-bound, so N nodes genuinely multiply throughput -- and on NCI `copyq`, +# already hpc.DEFAULT_TRANSFER_QUEUE, is the queue with outbound network access. +# A local destination is bound by the filesystem both ends share, where a second +# node adds contention rather than bandwidth; `-j` is the knob there. +# +# What makes this cheap to build is that the submitter has to resolve every row +# anyway, because LPT bin-packing needs the byte sizes and `list_source_files` +# is where sizes come from. Having paid for that, it ships the resolved +# (md5, size, key) tuples in the manifest -- so no worker ever runs `dvc list`, +# which is what keeps N nodes off DVC's SQLite state db for the shared clone. +# Each worker opens the clone once, read-only, for one thing only: the source +# remote's filesystem and credentials out of its .dvc/config. + + +def _worker_tasks( + targets: Sequence[Tuple[str, Optional[str]]], + resolved: Sequence[Tuple[Optional[List[Dict[str, object]]], Optional[str]]], +) -> Tuple[List[Dict[str, object]], Dict[int, int]]: + """Flatten resolved rows into JSON-serialisable per-file tasks. + + The partition unit is a *file*, not a row. A row keeps its identity through + the ``row`` tag, so reporting stays per-row either way, and per-file + balances properly in the two cases per-row cannot: a single row holding + thousands of files, and rows whose sizes differ by an order of magnitude -- + which fastq sample directories reliably do. + + Returns the task list, and a row index -> file count map so collation can + tell "this row's files failed" from "nothing came back for this row". + """ + tasks: List[Dict[str, object]] = [] + counts: Dict[int, int] = {} + for i, ((_, name), (entries, _)) in enumerate(zip(targets, resolved)): + if entries is None: + continue + counts[i] = len(entries) + for entry in entries: + tasks.append({ + 'row': i, + 'name': name, + 'relpath': str(entry['relpath']), + 'md5': str(entry['md5']), + 'size': int(entry.get('size') or 0), + }) + return tasks, counts + + +def _write_worker_results( + manifest_dir: Path, + worker_id: int, + results: List[Dict[str, object]], +) -> Path: + """Record one worker's per-file outcomes for the submitter to collate. + + Written even for an empty partition, so that a *missing* file means "this + worker did not finish" rather than "this worker had nothing to do". + """ + path = manifest_dir / f'result_{worker_id}.json' + with open(path, 'w') as f: + json.dump({'results': results}, f) + return path + + +def _collate_worker_results( + manifest_dir: Path, + targets: Sequence[Tuple[str, Optional[str]]], + resolved: Sequence[Tuple[Optional[List[Dict[str, object]]], Optional[str]]], + counts: Dict[int, int], + root: 'get_dest.S3Dest', +) -> List[Tuple[str, bool, str]]: + """Rebuild per-row outcomes, in CSV order, from the workers' result files. + + Because the partition unit is a file, one row's files may be spread across + several workers; the ``row`` tag each task carries is what puts them back + together. + + A row whose files do not all come back is reported as a failure even when + every result that *did* arrive succeeded. A worker killed mid-partition -- + walltime, most likely -- would otherwise read as a clean run of a short row. + """ + per_row: Dict[int, List[Tuple[str, bool, str]]] = {} + for part in sorted(manifest_dir.glob('result_*.json')): + try: + with open(part) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + # A truncated result file is itself a worker that did not finish, + # and the per-row count check below is what reports it. + continue + for item in data.get('results', []): + per_row.setdefault(item.get('row'), []).append(( + str(item.get('relpath')), + bool(item.get('ok')), + str(item.get('message')), + )) + + results: List[Tuple[str, bool, str]] = [] + for i, ((row_path, name), (entries, error)) in enumerate( + zip(targets, resolved) + ): + if entries is None: + results.append((row_path or '(empty)', False, error)) + continue + placed = per_row.get(i, []) + written = sum(1 for _, ok, _ in placed if ok) + failed = len(placed) - written + missing = counts.get(i, 0) - len(placed) + if failed or missing: + note = f'{written} uploaded, {failed} failed' + if missing: + note += f', {missing} unreported (a worker did not finish)' + results.append((row_path, False, note)) + else: + dest = root.under(name) if name else root + plural = '' if written == 1 else 's' + results.append((row_path, True, f'{written} file{plural} -> {dest}')) + return results + + +def distributed_upload_to_s3( + repository: str, + targets: Sequence[Tuple[str, Optional[str]]], + root: 'get_dest.S3Dest', + num_workers: int, + owner: Optional[str] = None, + rev: Optional[str] = None, + jobs: int = 8, + force: bool = False, + resume: bool = False, + check: bool = False, + chunk_size: int = get_dest.DEFAULT_CHUNK_SIZE, + refresh: bool = True, + path_col: str = 'path', + qxub_args: Optional[List[str]] = None, + wait: bool = True, + verbose: bool = False, +) -> Tuple[List[str], Optional[Path], Optional[List[Tuple[str, bool, str]]]]: + """Spread a stream-to-S3 transfer across ``num_workers`` compute nodes. + + Args: + repository: Repository name, alias, or URL. + targets: ``(source path, destination name)`` pairs, one per row. The + name is the sub-prefix under *root* the row lands in, or None to + use *root* itself. + root: The destination prefix, already preflighted by the caller. + num_workers: Number of qxub jobs to submit. + owner: Owner override for short repository names. + rev: Git revision to resolve against. + jobs: Streaming threads *within* each worker. + force: Overwrite existing objects. + resume: Skip objects already present. + check: Verify existing objects against their recorded checksum. + chunk_size: Streaming chunk size in bytes. + refresh: Refresh the cached clone before resolving. + path_col: CSV column name, for error messages only. + qxub_args: Additional arguments for qxub exec. + wait: Wait for the jobs and collate their results. + verbose: Print progress. + + Returns: + ``(job_ids, manifest_dir, results)``. *results* is the per-row report + when *wait* is True, and None when the jobs were left running. + + Raises: + GetError: If qxub is unavailable, nothing resolves, or no job submits. + """ + try: + hpc.require_qxub() + except hpc.HPCError as e: + raise GetError(str(e)) + + clone = tmp_mod.clone_repo( + repository, owner=owner, refresh=refresh, verbose=verbose, rev=rev + ) + + resolved = _resolve_rows( + clone, [row_path for row_path, _ in targets], rev, jobs, path_col + ) + tasks, counts = _worker_tasks(targets, resolved) + if not tasks: + problems = [error for _, error in resolved if error] + detail = f" First error: {problems[0]}" if problems else "" + raise GetError(f"No row resolved to a tracked file.{detail}") + + partitions = hpc.partition_by_size( + tasks, + num_workers, + weight=lambda task: task['size'], + # (row, relpath) is unique across the manifest and independent of the + # order the rows happened to resolve in, so the same CSV always + # partitions the same way. + tiebreak=lambda task: (task['row'], task['relpath']), + verbose=verbose, + ) + + job_id = str(uuid.uuid4())[:8] + manifest_dir = hpc.save_manifest( + {'files': tasks, 'repo_root': str(Path.cwd())}, + partitions, + job_id, + operation='get', + metadata={ + # Resolved on the submitting node and passed through, not + # re-derived per worker. resolve_link_types() reads `dvc config` + # from the cwd and preflight() resolves a credential chain; both + # could answer differently on a compute node, and a transfer that + # silently changes its mind halfway is worse than one that fails. + 'clone': str(clone), + 'rev': rev, + 'root_url': root.url, + 'dest_config': dataclasses.asdict(root.config), + 'chunk_size': chunk_size, + 'jobs': jobs, + 'force': force, + 'resume': resume, + 'check': check, + }, + ) + + if verbose: + print(f"Manifest saved to {manifest_dir}") + + try: + job_ids = hpc.submit_workers( + manifest_dir, + num_workers, + operation='get', + qxub_args=qxub_args, + verbose=verbose, + ) + except hpc.HPCError as e: + raise GetError(str(e)) + + if not job_ids: + raise GetError( + f"No worker jobs were submitted. The manifest is at {manifest_dir} " + f"if you want to inspect the partitioning." + ) + + if not wait: + return job_ids, manifest_dir, None + + try: + succeeded = hpc.monitor_jobs(job_ids, verbose=verbose) + except hpc.HPCError as e: + raise GetError(str(e)) + if not succeeded and verbose: + print("Some worker jobs reported failure; collating what landed.") + + return job_ids, manifest_dir, _collate_worker_results( + manifest_dir, targets, resolved, counts, root + ) + + +def worker_upload_to_s3( + manifest_dir: Path, + worker_id: int, + jobs: Optional[int] = None, + verbose: bool = False, +) -> Tuple[int, int]: + """Stream one worker's partition into S3. Called inside a submitted job. + + The partition arrives already resolved -- md5, size and destination for + every file -- so this never runs ``dvc list``, and the only thing it needs + the clone for is the source remote's filesystem and credentials. + + Args: + manifest_dir: Manifest directory written by the submitter. + worker_id: Worker index. + jobs: Streaming threads. None takes the value the submitter chose, + which is the authoritative one. + verbose: Print progress. + + Returns: + (uploaded_count, failed_count) for this partition. + """ + metadata, tasks = hpc.load_worker_partition(manifest_dir, worker_id) + + root = get_dest.S3Dest( + metadata['root_url'], + get_dest.S3DestConfig(**(metadata.get('dest_config') or {})), + ) + + # Every worker preflights for itself rather than trusting the submitter's + # check. --dest-account-id exists for unattended runs and a worker is the + # most unattended thing here: a credential chain that resolves differently + # on a compute node than on the login node is exactly what it must catch. + identity = root.preflight() + if verbose: + rows = len({task.get('row') for task in tasks}) + print(f"Worker {worker_id}: {len(tasks)} object(s) from {rows} " + f"row{'' if rows == 1 else 's'} under {root}") + print(f"Writing as: {identity}") + + if not tasks: + _write_worker_results(manifest_dir, worker_id, []) + return 0, 0 + + jobs = jobs or int(metadata.get('jobs') or 8) + chunk_size = int(metadata.get('chunk_size') or get_dest.DEFAULT_CHUNK_SIZE) + + # Build the filesystem once up front: under() copies it into each child, so + # touching it here means every row shares one client rather than each + # relying on fsspec's instance cache to hand back the same one. + _ = root.fs + dests: Dict[Optional[str], 'get_dest.S3Dest'] = {} + for task in tasks: + name = task.get('name') + if name not in dests: + dests[name] = root.under(name) if name else root + + repo, odb = _open_source_odb(Path(metadata['clone'])) + try: + placed = _upload_all( + [ + ({'relpath': task['relpath'], 'md5': task['md5']}, + dests[task.get('name')]) + for task in tasks + ], + odb, + jobs, + bool(metadata.get('force')), + bool(metadata.get('resume')), + bool(metadata.get('check')), + chunk_size, + ) + finally: + repo.close() + + _write_worker_results(manifest_dir, worker_id, [ + {'row': task['row'], 'relpath': relpath, 'ok': ok, 'message': message} + for task, (relpath, ok, message) in zip(tasks, placed) + ]) + return _tally(placed, verbose) + + def _prepare_s3_dest( out: str, name: Optional[str], @@ -824,39 +1220,22 @@ def get_from_csv( link_types = resolve_link_types(link) - # Phase 1: resolve every row to a file listing. Each row costs a `dvc list` - # subprocess, so run them concurrently -- on a long manifest the resolves - # alone are otherwise a serial prologue to all the real work. - def resolve(target): - row_path, row_out = target - if not row_path: - return None, None, f'Missing {path_col}' - try: - entries = list_source_files(clone, row_path, rev=rev) - except GetError as e: - return None, None, str(e) - return _dest_for(row_out, row_path), entries, None - - # ...but resolve the first row on its own first. DVC keeps a SQLite state db - # inside the clone, and on a freshly-cloned repo it does not exist yet; fan - # out immediately and the workers race to create it, which fails all but one - # with "database is locked". One serial call builds it, after which - # concurrent readers are fine. - resolved = [] - if targets: - resolved.append(resolve(targets[0])) - if len(targets) > 1: - with ThreadPoolExecutor(max_workers=max(1, jobs)) as pool: - resolved.extend(pool.map(resolve, targets[1:])) + # Phase 1: resolve every row to a file listing. + resolved = _resolve_rows( + clone, [row_path for row_path, _ in targets], rev, jobs, path_col + ) # Phase 2: place every file from every row through a single pool, so the # worker budget goes to the transfer as a whole rather than being re-divided # per row. Spans let each row's outcome be recovered from the flat results. tasks: List[Tuple[Path, Dict[str, object]]] = [] spans: Dict[int, Tuple[Path, int, int]] = {} - for i, (dest_root, entries, _) in enumerate(resolved): + for i, ((row_path, row_out), (entries, _)) in enumerate( + zip(targets, resolved) + ): if entries is None: continue + dest_root = _dest_for(row_out, row_path) start = len(tasks) tasks.extend((dest_root, entry) for entry in entries) spans[i] = (dest_root, start, len(tasks)) @@ -870,7 +1249,7 @@ def resolve(target): # Phase 3: attribute the flat results back to their rows, in CSV order. results: List[Tuple[str, bool, str]] = [] - for i, (target, (_, entries, error)) in enumerate(zip(targets, resolved)): + for i, (target, (entries, error)) in enumerate(zip(targets, resolved)): row_path = target[0] if entries is None: results.append((row_path or '(empty)', False, error)) diff --git a/dt/hpc.py b/dt/hpc.py index ea1963a..8251341 100644 --- a/dt/hpc.py +++ b/dt/hpc.py @@ -249,52 +249,136 @@ def get_prefixes_for_worker(worker_id: int, num_workers: int) -> Set[str]: return prefixes +def partition_by_size( + items: List[Any], + num_workers: int, + weight: Callable[[Any], int], + tiebreak: Callable[[Any], Any], + verbose: bool = False, +) -> Dict[int, List[Any]]: + """Split *items* across workers, balancing total weight per worker. + + Greedy longest-processing-time (LPT) bin-packing: items are sorted + largest-first and each is assigned to the currently least-loaded worker. + Wave wall-clock then approaches + ``max(total_weight / num_workers, largest_single_item) / per-worker-rate`` + instead of being dominated by whichever worker an arbitrary split happened + to hand the big items (issue #138). + + This lives here, rather than in any one caller, because it is the answer to + a problem *qxub workers specifically* have: they are separate processes on + separate machines and cannot steal work from each other, so the split has to + be decided up front. A ``ThreadPoolExecutor`` hands tasks out from a shared + queue and so load-balances at runtime -- applying LPT there would be + strictly worse, freezing an assignment the queue would otherwise self-correct + (issue #172). Reach for this only when the consumers are separate processes. + + Each item is assigned to exactly one worker, so partitions are disjoint and + workers never contend. + + Args: + items: Work items to distribute. Opaque here; only *weight* and + *tiebreak* look inside them. + num_workers: Number of workers to split across. + weight: Cost of an item, in whatever unit the caller balances (bytes, + for both current callers). + tiebreak: Total order on equal-weight items. Required, not optional: + without it the assignment depends on the input order, and a + partitioning you cannot reproduce is one you cannot debug. + verbose: Log the resulting per-worker balance. + + Returns: + Dict mapping worker_id to its list of items. + + Raises: + ValueError: If *num_workers* is less than 1. + """ + if num_workers < 1: + raise ValueError(f"num_workers must be at least 1, got {num_workers}") + + partitions: Dict[int, List[Any]] = {i: [] for i in range(num_workers)} + loads = [0] * num_workers + + for item in sorted(items, key=lambda it: (weight(it), tiebreak(it)), + reverse=True): + # Assign to the least-loaded worker; tie-break by lowest worker id. + worker_id = min(range(num_workers), key=lambda w: (loads[w], w)) + partitions[worker_id].append(item) + loads[worker_id] += weight(item) + + if verbose: + active = [load for load in loads if load > 0] + if active: + print( + f"Partition byte balance across {len(active)} active worker(s): " + f"min={utils.format_size(min(active))}, " + f"max={utils.format_size(max(active))}, " + f"total={utils.format_size(sum(loads))}" + ) + + return partitions + + def save_manifest( manifest: Dict[str, Any], - partitions: Dict[int, List[str]], + partitions: Dict[int, List[Any]], job_id: str, operation: str, + metadata: Optional[Dict[str, Any]] = None, ) -> Path: """Save manifest and partitions to disk. - + Args: manifest: Original manifest with metadata - partitions: Worker partitions (worker_id -> list of file hashes) + partitions: Worker partitions (worker_id -> list of work items) job_id: Unique job identifier - operation: Transfer operation name ('push' or 'pull') - + operation: Transfer operation name ('push', 'pull' or 'get') + metadata: Extra key/values to record in ``manifest.json`` alongside the + standard fields. Everything a worker needs beyond its own partition + travels this way, because the worker is a fresh process on another + machine with no share of the submitter's state: ``push`` needs only + a remote name and a repo root, but ``get`` also has to be handed a + destination, credentials config and the flags the run was invoked + with. Read back by :func:`load_worker_partition`. + Returns: Path to the manifest directory """ manifest_dir = get_transfer_dir(operation) / job_id manifest_dir.mkdir(parents=True, exist_ok=True) - + # Save metadata + payload = { + 'remote': manifest.get('remote'), + 'repo_root': manifest.get('repo_root'), + 'total_files': len(manifest.get('files', [])), + 'num_workers': len(partitions), + } + payload.update(metadata or {}) with open(manifest_dir / 'manifest.json', 'w') as f: - json.dump({ - 'remote': manifest.get('remote'), - 'repo_root': manifest.get('repo_root'), - 'total_files': len(manifest.get('files', [])), - 'num_workers': len(partitions), - }, f, indent=2) - + json.dump(payload, f, indent=2) + # Save each worker's partition for worker_id, files in partitions.items(): with open(manifest_dir / f'worker_{worker_id}.json', 'w') as f: json.dump({'files': files}, f) - + return manifest_dir -def load_worker_partition(manifest_dir: Path, worker_id: int) -> Tuple[Dict, List[str]]: - """Load manifest metadata and worker's file partition. - +def load_worker_partition( + manifest_dir: Path, worker_id: int +) -> Tuple[Dict, List[Any]]: + """Load manifest metadata and worker's partition. + Args: manifest_dir: Path to manifest directory worker_id: Worker index - + Returns: - Tuple of (metadata dict, list of file hashes) + Tuple of (metadata dict, list of work items). The items are whatever + the submitter put in the partition -- hashes for ``push``, per-file + task dicts for ``get``. """ with open(manifest_dir / 'manifest.json') as f: metadata = json.load(f) diff --git a/dt/push.py b/dt/push.py index 433a8b0..5306ef7 100644 --- a/dt/push.py +++ b/dt/push.py @@ -321,15 +321,12 @@ def partition_manifest( ) -> Dict[int, List[str]]: """Partition manifest files across workers, balancing bytes per worker. - Uses greedy longest-processing-time (LPT) bin-packing: files are sorted - largest-first and each is assigned to the currently least-loaded worker. - This balances the total bytes per worker, so wave wall-clock approaches - ``max(total_bytes / num_workers, largest_single_file) / per-worker-rate`` - instead of being dominated by whichever worker a hash-prefix split happened - to hand the big files (see issue #138). - - The key partitioning invariant is preserved: each file is assigned to - exactly one worker, so partitions are disjoint and workers never contend. + The bin-packing itself is :func:`dt.hpc.partition_by_size` -- it is not + push-specific, and ``dt get --workers`` balances its own work the same way + (issue #172). What stays here is the part that genuinely is push-specific: + a file is identified by its hash, and its size therefore has to be read out + of the local cache. + Assignment is deterministic given the file sizes (ties broken by hash, then by lowest worker id). @@ -348,28 +345,13 @@ def partition_manifest( if sizes is None: sizes = get_file_sizes(files) - partitions: Dict[int, List[str]] = {i: [] for i in range(num_workers)} - loads = [0] * num_workers - - # Largest files first; tie-break by hash so ordering is reproducible - # regardless of the manifest's file order. - for file_hash in sorted(files, key=lambda h: (sizes.get(h, 0), h), reverse=True): - # Assign to the least-loaded worker; tie-break by lowest worker id. - worker_id = min(range(num_workers), key=lambda w: (loads[w], w)) - partitions[worker_id].append(file_hash) - loads[worker_id] += sizes.get(file_hash, 0) - - if verbose: - active = [load for load in loads if load > 0] - if active: - print( - f"Partition byte balance across {len(active)} active worker(s): " - f"min={utils.format_size(min(active))}, " - f"max={utils.format_size(max(active))}, " - f"total={utils.format_size(sum(loads))}" - ) - - return partitions + return hpc.partition_by_size( + files, + num_workers, + weight=lambda h: sizes.get(h, 0), + tiebreak=lambda h: h, + verbose=verbose, + ) def _missing_from_remote( diff --git a/pyproject.toml b/pyproject.toml index 13ac11b..84b3126 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "dvc-tools" -version = "0.24.0" +version = "0.25.0" description = "Convenient tools for working with DVC in HPC environments" readme = "README.md" requires-python = ">=3.8" diff --git a/tests/unit/test_get_workers.py b/tests/unit/test_get_workers.py new file mode 100644 index 0000000..8b38662 --- /dev/null +++ b/tests/unit/test_get_workers.py @@ -0,0 +1,671 @@ +"""Unit tests for ``dt get --workers``: distributed streaming to S3 via qxub. + +No qxub and no network. The interesting behaviour is the *round trip* -- the +submitter resolves and partitions, each worker reads its partition back off +disk and reports, and the submitter reassembles per-row outcomes from those +reports. So the central test here replaces ``hpc.submit_workers`` with a stub +that runs ``worker_upload_to_s3`` inline, which exercises the manifest as a +real interface rather than asserting on how it happens to be spelled. + +What this cannot cover is anything about PBS: whether the queue has outbound +network access, whether the manifest directory is visible from a compute node, +or whether more nodes actually move more bytes. Those need a real submission. +""" + +import json +from pathlib import Path +from unittest.mock import patch + +import pytest + +from dt import get as get_mod +from dt import get_dest +from dt import hpc +from dt.errors import GetError + +# The S3 fakes are the ones the single-node tests already run against, so a +# behaviour change there fails here too rather than passing against a second, +# drifted copy. +from tests.unit.test_get_s3 import FakeODB, FakeS3FS, md5_of + + +# ============================================================================= +# Fakes +# ============================================================================= + +class FakeRepo: + """Stands in for the DVC Repo opened against the source clone.""" + + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + +@pytest.fixture +def fs(): + return FakeS3FS() + + +@pytest.fixture +def source(fs): + """A source remote holding two files per sample, of differing sizes.""" + blobs = { + 'AF013-A/R1.fq': b'a' * 400, + 'AF013-A/R2.fq': b'b' * 100, + 'AF013-B/R1.fq': b'c' * 50, + } + objects = {md5_of(data): data for data in blobs.values()} + odb = FakeODB(fs, objects) + entries = { + 'data/fq/AF013-A': [ + {'relpath': 'R1.fq', 'md5': md5_of(blobs['AF013-A/R1.fq']), 'size': 400}, + {'relpath': 'R2.fq', 'md5': md5_of(blobs['AF013-A/R2.fq']), 'size': 100}, + ], + 'data/fq/AF013-B': [ + {'relpath': 'R1.fq', 'md5': md5_of(blobs['AF013-B/R1.fq']), 'size': 50}, + ], + } + return odb, entries, blobs + + +def s3_patches(fs): + """Patch the destination's filesystem and preflight, leaving addressing real.""" + return ( + patch.object(get_dest.S3Dest, 'fs', property(lambda self: fs)), + patch.object(get_dest.S3Dest, 'preflight', + lambda self: 'account 000000000000 as arn:aws:iam::test'), + ) + + +def run_workers_inline(manifest_dir, num_workers, operation, **kwargs): + """Stand-in for hpc.submit_workers that runs each worker in-process.""" + job_ids = [] + for worker_id in range(num_workers): + if not (manifest_dir / f'worker_{worker_id}.json').exists(): + continue + _, tasks = hpc.load_worker_partition(manifest_dir, worker_id) + if not tasks: + continue + get_mod.worker_upload_to_s3(manifest_dir, worker_id) + job_ids.append(f'{worker_id}.fake-pbs') + return job_ids + + +# ============================================================================= +# hpc.partition_by_size -- the bin-packing itself +# ============================================================================= + +class TestPartitionBySize: + + def test_balances_bytes_rather_than_counts(self): + """One big item plus many small ones must not all land together.""" + items = [{'n': 'big', 'size': 1000}] + [ + {'n': str(i), 'size': 100} for i in range(10) + ] + parts = hpc.partition_by_size( + items, 2, weight=lambda i: i['size'], tiebreak=lambda i: i['n']) + + loads = [sum(i['size'] for i in v) for v in parts.values()] + # 2000 bytes over 2 workers, and the 1000-byte item is indivisible. + assert sorted(loads) == [1000, 1000] + + def test_partitions_are_disjoint_and_complete(self): + items = [{'n': str(i), 'size': i} for i in range(20)] + parts = hpc.partition_by_size( + items, 4, weight=lambda i: i['size'], tiebreak=lambda i: i['n']) + + names = [i['n'] for v in parts.values() for i in v] + assert sorted(names) == sorted(i['n'] for i in items) + assert len(names) == len(set(names)) + + def test_assignment_is_independent_of_input_order(self): + items = [{'n': str(i), 'size': (i % 3) * 10} for i in range(12)] + args = dict(weight=lambda i: i['size'], tiebreak=lambda i: i['n']) + + forward = hpc.partition_by_size(items, 3, **args) + backward = hpc.partition_by_size(list(reversed(items)), 3, **args) + + assert forward == backward + + def test_every_worker_gets_a_key_even_with_nothing_to_do(self): + parts = hpc.partition_by_size( + [{'n': 'only', 'size': 1}], 4, + weight=lambda i: i['size'], tiebreak=lambda i: i['n']) + + assert set(parts) == {0, 1, 2, 3} + assert sum(len(v) for v in parts.values()) == 1 + + def test_rejects_zero_workers(self): + with pytest.raises(ValueError, match='at least 1'): + hpc.partition_by_size( + [], 0, weight=lambda i: 0, tiebreak=lambda i: '') + + +class TestSaveManifestMetadata: + + def test_metadata_reaches_the_worker(self, tmp_path, monkeypatch): + """Whatever the submitter records must come back out the other side.""" + monkeypatch.chdir(tmp_path) + manifest_dir = hpc.save_manifest( + {'files': ['x']}, {0: ['x']}, 'job1', 'get', + metadata={'root_url': 's3://b/p', 'jobs': 4}, + ) + + metadata, files = hpc.load_worker_partition(manifest_dir, 0) + + assert metadata['root_url'] == 's3://b/p' + assert metadata['jobs'] == 4 + assert files == ['x'] + + def test_standard_fields_survive_without_metadata(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + manifest_dir = hpc.save_manifest( + {'files': ['a', 'b'], 'remote': 'r2', 'repo_root': '/repo'}, + {0: ['a'], 1: ['b']}, 'job2', 'push', + ) + + metadata, _ = hpc.load_worker_partition(manifest_dir, 1) + + assert metadata['remote'] == 'r2' + assert metadata['repo_root'] == '/repo' + assert metadata['total_files'] == 2 + assert metadata['num_workers'] == 2 + + +# ============================================================================= +# Task construction and collation +# ============================================================================= + +class TestWorkerTasks: + + def test_tags_every_file_with_its_row(self): + targets = [('data/a', 'a'), ('data/b', 'b')] + resolved = [ + ([{'relpath': 'f1', 'md5': 'x' * 32, 'size': 10}, + {'relpath': 'f2', 'md5': 'y' * 32, 'size': 20}], None), + ([{'relpath': 'f3', 'md5': 'z' * 32, 'size': 30}], None), + ] + + tasks, counts = get_mod._worker_tasks(targets, resolved) + + assert [t['row'] for t in tasks] == [0, 0, 1] + assert [t['name'] for t in tasks] == ['a', 'a', 'b'] + assert counts == {0: 2, 1: 1} + + def test_unresolved_rows_contribute_no_tasks_and_no_count(self): + targets = [('data/a', 'a'), ('', None), ('data/c', 'c')] + resolved = [ + ([{'relpath': 'f1', 'md5': 'x' * 32, 'size': 1}], None), + (None, 'Missing path'), + (None, 'Not found in the source repository: data/c'), + ] + + tasks, counts = get_mod._worker_tasks(targets, resolved) + + assert [t['row'] for t in tasks] == [0] + # No count for a row that never resolved -- collation must not read + # that as "0 of 0 files arrived, so this row succeeded". + assert counts == {0: 1} + + +class TestCollateWorkerResults: + + root = None + + def _root(self, fs): + dest = get_dest.S3Dest('s3://bucket/fastqs') + dest._fs = fs + return dest + + def _write(self, manifest_dir, worker_id, results): + (manifest_dir / f'result_{worker_id}.json').write_text( + json.dumps({'results': results})) + + def test_rows_reported_in_order_though_files_span_workers(self, tmp_path, fs): + targets = [('data/a', 'a'), ('data/b', 'b')] + resolved = [([{}, {}], None), ([{}], None)] + counts = {0: 2, 1: 1} + # Row 0's two files were split across two workers. + self._write(tmp_path, 0, [ + {'row': 0, 'relpath': 'f1', 'ok': True, 'message': 'uploaded'}, + {'row': 1, 'relpath': 'f3', 'ok': True, 'message': 'uploaded'}, + ]) + self._write(tmp_path, 1, [ + {'row': 0, 'relpath': 'f2', 'ok': True, 'message': 'uploaded'}, + ]) + + results = get_mod._collate_worker_results( + tmp_path, targets, resolved, counts, self._root(fs)) + + assert [r[0] for r in results] == ['data/a', 'data/b'] + assert all(ok for _, ok, _ in results) + assert results[0][2] == '2 files -> s3://bucket/fastqs/a' + assert results[1][2] == '1 file -> s3://bucket/fastqs/b' + + def test_unresolved_row_keeps_its_own_error(self, tmp_path, fs): + targets = [('data/a', 'a'), ('data/gone', 'gone')] + resolved = [([{}], None), (None, 'Not found in the source repository')] + self._write(tmp_path, 0, [ + {'row': 0, 'relpath': 'f1', 'ok': True, 'message': 'uploaded'}, + ]) + + results = get_mod._collate_worker_results( + tmp_path, targets, resolved, {0: 1}, self._root(fs)) + + assert results[1] == ('data/gone', False, + 'Not found in the source repository') + + def test_a_worker_that_never_reported_fails_its_rows(self, tmp_path, fs): + """The failure mode that matters: a job killed on walltime. + + Every result that arrived succeeded, so a naive tally would call this a + clean run of a one-file row. + """ + targets = [('data/a', 'a')] + resolved = [([{}, {}, {}], None)] + self._write(tmp_path, 0, [ + {'row': 0, 'relpath': 'f1', 'ok': True, 'message': 'uploaded'}, + ]) + + results = get_mod._collate_worker_results( + tmp_path, targets, resolved, {0: 3}, self._root(fs)) + + assert results[0][1] is False + assert '2 unreported' in results[0][2] + + def test_no_result_files_at_all_is_a_failure(self, tmp_path, fs): + results = get_mod._collate_worker_results( + tmp_path, [('data/a', 'a')], [([{}], None)], {0: 1}, + self._root(fs)) + + assert results[0][1] is False + assert '1 unreported' in results[0][2] + + def test_reported_failures_are_counted(self, tmp_path, fs): + self._write(tmp_path, 0, [ + {'row': 0, 'relpath': 'f1', 'ok': True, 'message': 'uploaded'}, + {'row': 0, 'relpath': 'f2', 'ok': False, + 'message': 'object abc not on the source remote'}, + ]) + + results = get_mod._collate_worker_results( + tmp_path, [('data/a', 'a')], [([{}, {}], None)], {0: 2}, + self._root(fs)) + + assert results[0] == ('data/a', False, '1 uploaded, 1 failed') + + def test_a_truncated_result_file_reads_as_unfinished(self, tmp_path, fs): + (tmp_path / 'result_0.json').write_text('{"results": [{"row": 0,') + + results = get_mod._collate_worker_results( + tmp_path, [('data/a', 'a')], [([{}], None)], {0: 1}, + self._root(fs)) + + assert results[0][1] is False + assert 'unreported' in results[0][2] + + +# ============================================================================= +# The worker +# ============================================================================= + +class TestWorkerUploadToS3: + + def _manifest(self, tmp_path, source, tasks, **metadata): + odb, entries, _ = source + base = { + 'clone': str(tmp_path / 'clone'), + 'rev': None, + 'root_url': 's3://bucket/fastqs', + 'dest_config': {'profile': None, 'endpoint_url': None, + 'region': None, 'account_id': None}, + 'chunk_size': 64, + 'jobs': 2, + 'force': False, 'resume': False, 'check': False, + } + base.update(metadata) + return hpc.save_manifest( + {'files': tasks}, {0: tasks}, 'jobw', 'get', metadata=base) + + def test_uploads_its_partition_and_records_every_file( + self, tmp_path, monkeypatch, fs, source + ): + monkeypatch.chdir(tmp_path) + odb, entries, blobs = source + tasks = [ + {'row': 0, 'name': 'AF013-A', 'relpath': 'R1.fq', + 'md5': md5_of(blobs['AF013-A/R1.fq']), 'size': 400}, + {'row': 1, 'name': 'AF013-B', 'relpath': 'R1.fq', + 'md5': md5_of(blobs['AF013-B/R1.fq']), 'size': 50}, + ] + manifest_dir = self._manifest(tmp_path, source, tasks) + repo = FakeRepo() + + p1, p2 = s3_patches(fs) + with p1, p2, patch.object(get_mod, '_open_source_odb', + return_value=(repo, odb)): + written, failed = get_mod.worker_upload_to_s3(manifest_dir, 0) + + assert (written, failed) == (2, 0) + # The objects landed under the right per-row prefix... + assert fs.objects['s3://bucket/fastqs/AF013-A/R1.fq'] == \ + blobs['AF013-A/R1.fq'] + assert fs.objects['s3://bucket/fastqs/AF013-B/R1.fq'] == \ + blobs['AF013-B/R1.fq'] + # ...and the row tags survived into the result file. + recorded = json.loads( + (manifest_dir / 'result_0.json').read_text())['results'] + assert sorted(r['row'] for r in recorded) == [0, 1] + assert all(r['ok'] for r in recorded) + assert repo.closed + + def test_never_lists_the_source_repository( + self, tmp_path, monkeypatch, fs, source + ): + """The whole point of shipping resolved tasks. + + A worker that ran ``dvc list`` would put N compute nodes onto the + shared clone's SQLite state db. + """ + monkeypatch.chdir(tmp_path) + odb, _, blobs = source + tasks = [{'row': 0, 'name': 'AF013-A', 'relpath': 'R1.fq', + 'md5': md5_of(blobs['AF013-A/R1.fq']), 'size': 400}] + manifest_dir = self._manifest(tmp_path, source, tasks) + + p1, p2 = s3_patches(fs) + with p1, p2, \ + patch.object(get_mod, '_open_source_odb', + return_value=(FakeRepo(), odb)), \ + patch.object(get_mod, 'list_source_files') as listing: + get_mod.worker_upload_to_s3(manifest_dir, 0) + + listing.assert_not_called() + + def test_an_empty_partition_still_writes_a_result_file( + self, tmp_path, monkeypatch, fs, source + ): + """Absence of a result file has to mean "did not finish".""" + monkeypatch.chdir(tmp_path) + manifest_dir = self._manifest(tmp_path, source, []) + + p1, p2 = s3_patches(fs) + with p1, p2: + written, failed = get_mod.worker_upload_to_s3(manifest_dir, 0) + + assert (written, failed) == (0, 0) + assert json.loads( + (manifest_dir / 'result_0.json').read_text()) == {'results': []} + + def test_preflights_for_itself(self, tmp_path, monkeypatch, fs, source): + """--dest-account-id must be enforced on the node doing the writing.""" + monkeypatch.chdir(tmp_path) + odb, _, blobs = source + tasks = [{'row': 0, 'name': 'AF013-A', 'relpath': 'R1.fq', + 'md5': md5_of(blobs['AF013-A/R1.fq']), 'size': 400}] + manifest_dir = self._manifest( + tmp_path, source, + tasks, + dest_config={'profile': None, 'endpoint_url': None, + 'region': None, 'account_id': '999999999999'}, + ) + + with patch.object(get_dest.S3Dest, 'fs', property(lambda self: fs)), \ + patch.object(get_dest.S3Dest, 'preflight', + side_effect=GetError('Destination account ' + 'mismatch')), \ + patch.object(get_mod, '_open_source_odb', + return_value=(FakeRepo(), odb)): + with pytest.raises(GetError, match='account mismatch'): + get_mod.worker_upload_to_s3(manifest_dir, 0) + + def test_submitter_jobs_value_is_authoritative( + self, tmp_path, monkeypatch, fs, source + ): + monkeypatch.chdir(tmp_path) + odb, _, blobs = source + tasks = [{'row': 0, 'name': 'A', 'relpath': 'R1.fq', + 'md5': md5_of(blobs['AF013-A/R1.fq']), 'size': 400}] + manifest_dir = self._manifest(tmp_path, source, tasks, jobs=5) + + p1, p2 = s3_patches(fs) + with p1, p2, \ + patch.object(get_mod, '_open_source_odb', + return_value=(FakeRepo(), odb)), \ + patch.object(get_mod, '_upload_all', + return_value=[('R1.fq', True, 'uploaded')]) as up: + get_mod.worker_upload_to_s3(manifest_dir, 0) + + assert up.call_args[0][2] == 5 + + +# ============================================================================= +# The submitter, end to end through a real manifest +# ============================================================================= + +class TestDistributedUploadToS3: + + def _root(self, fs): + dest = get_dest.S3Dest('s3://bucket/fastqs') + dest._fs = fs + return dest + + def _run(self, tmp_path, fs, source, targets, num_workers=2, + submit=run_workers_inline, monitor=True): + odb, entries, _ = source + + def listing(clone, path, rev=None): + if path not in entries: + raise GetError(f'Not found in the source repository: {path}') + return entries[path] + + p1, p2 = s3_patches(fs) + with p1, p2, \ + patch.object(get_mod.tmp_mod, 'clone_repo', + return_value=tmp_path / 'clone'), \ + patch.object(get_mod, 'list_source_files', side_effect=listing), \ + patch.object(get_mod, '_open_source_odb', + return_value=(FakeRepo(), odb)), \ + patch.object(hpc, 'require_qxub'), \ + patch.object(hpc, 'submit_workers', side_effect=submit), \ + patch.object(hpc, 'monitor_jobs', return_value=monitor): + return get_mod.distributed_upload_to_s3( + repository='myrepo', + targets=targets, + root=self._root(fs), + num_workers=num_workers, + jobs=2, + chunk_size=64, + ) + + def test_round_trip_places_every_object_and_reports_per_row( + self, tmp_path, monkeypatch, fs, source + ): + monkeypatch.chdir(tmp_path) + _, _, blobs = source + + job_ids, manifest_dir, results = self._run( + tmp_path, fs, source, + [('data/fq/AF013-A', 'AF013-A'), ('data/fq/AF013-B', 'AF013-B')], + ) + + assert len(job_ids) >= 1 + assert [r[0] for r in results] == \ + ['data/fq/AF013-A', 'data/fq/AF013-B'] + assert all(ok for _, ok, _ in results), results + assert set(fs.objects) >= { + 's3://bucket/fastqs/AF013-A/R1.fq', + 's3://bucket/fastqs/AF013-A/R2.fq', + 's3://bucket/fastqs/AF013-B/R1.fq', + } + for name, data in blobs.items(): + assert fs.objects[f's3://bucket/fastqs/{name}'] == data + + def test_partitions_by_size_not_by_row( + self, tmp_path, monkeypatch, fs, source + ): + """Row A holds 500 bytes and row B holds 50, over two workers. + + A per-row split would give one worker 10x the other's bytes; a per-file + split puts A's 400-byte file alone and the other two together. + """ + monkeypatch.chdir(tmp_path) + _, _, blobs = source + + _, manifest_dir, _ = self._run( + tmp_path, fs, source, + [('data/fq/AF013-A', 'AF013-A'), ('data/fq/AF013-B', 'AF013-B')], + ) + + loads = [] + for worker_id in (0, 1): + _, tasks = hpc.load_worker_partition(manifest_dir, worker_id) + loads.append(sum(t['size'] for t in tasks)) + assert sorted(loads) == [150, 400] + + def test_a_single_path_with_no_trailing_slash_lands_on_the_root( + self, tmp_path, monkeypatch, fs, source + ): + """``-o s3://b/p`` uses the prefix verbatim, so the row has no name. + + That None has to survive a JSON round trip through the manifest and + come back as "use the root itself" rather than a prefix literally + called "None". + """ + monkeypatch.chdir(tmp_path) + _, _, blobs = source + + _, _, results = self._run( + tmp_path, fs, source, [('data/fq/AF013-B', None)], + ) + + assert results == [('data/fq/AF013-B', True, + '1 file -> s3://bucket/fastqs')] + assert fs.objects['s3://bucket/fastqs/R1.fq'] == blobs['AF013-B/R1.fq'] + + def test_a_row_that_does_not_resolve_is_reported_not_fatal( + self, tmp_path, monkeypatch, fs, source + ): + monkeypatch.chdir(tmp_path) + + _, _, results = self._run( + tmp_path, fs, source, + [('data/fq/AF013-A', 'AF013-A'), ('data/fq/nope', 'nope')], + ) + + assert results[0][1] is True + assert results[1][1] is False + assert 'Not found' in results[1][2] + + def test_nothing_resolving_at_all_is_an_error( + self, tmp_path, monkeypatch, fs, source + ): + monkeypatch.chdir(tmp_path) + + with pytest.raises(GetError, match='No row resolved'): + self._run(tmp_path, fs, source, [('data/fq/nope', 'nope')]) + + def test_no_wait_returns_before_collating( + self, tmp_path, monkeypatch, fs, source + ): + monkeypatch.chdir(tmp_path) + odb, entries, _ = source + + p1, p2 = s3_patches(fs) + with p1, p2, \ + patch.object(get_mod.tmp_mod, 'clone_repo', + return_value=tmp_path / 'clone'), \ + patch.object(get_mod, 'list_source_files', + side_effect=lambda c, p, rev=None: entries[p]), \ + patch.object(hpc, 'require_qxub'), \ + patch.object(hpc, 'submit_workers', return_value=['1.pbs']), \ + patch.object(hpc, 'monitor_jobs') as monitor: + job_ids, manifest_dir, results = get_mod.distributed_upload_to_s3( + repository='myrepo', + targets=[('data/fq/AF013-A', 'AF013-A')], + root=self._root(fs), + num_workers=2, + wait=False, + ) + + assert job_ids == ['1.pbs'] + assert results is None + assert manifest_dir.exists() + monitor.assert_not_called() + + def test_failing_to_submit_anything_names_the_manifest( + self, tmp_path, monkeypatch, fs, source + ): + monkeypatch.chdir(tmp_path) + + with pytest.raises(GetError, match='No worker jobs were submitted'): + self._run(tmp_path, fs, source, + [('data/fq/AF013-A', 'AF013-A')], + submit=lambda *a, **k: []) + + def test_missing_qxub_is_a_get_error(self, tmp_path, monkeypatch, fs, source): + monkeypatch.chdir(tmp_path) + + with patch.object(hpc, 'require_qxub', + side_effect=hpc.HPCError('qxub not found')): + with pytest.raises(GetError, match='qxub not found'): + get_mod.distributed_upload_to_s3( + repository='myrepo', + targets=[('data/fq/AF013-A', 'AF013-A')], + root=self._root(fs), + num_workers=2, + ) + + def test_worker_results_survive_a_reported_job_failure( + self, tmp_path, monkeypatch, fs, source + ): + """A non-zero PBS exit must not discard what the workers did place.""" + monkeypatch.chdir(tmp_path) + + _, _, results = self._run( + tmp_path, fs, source, + [('data/fq/AF013-A', 'AF013-A')], + monitor=False, + ) + + assert results[0][1] is True + + def test_the_manifest_carries_no_secrets( + self, tmp_path, monkeypatch, fs, source + ): + """Profile name, endpoint and region travel; key material never does.""" + monkeypatch.chdir(tmp_path) + odb, entries, _ = source + root = get_dest.S3Dest( + 's3://bucket/fastqs', + get_dest.S3DestConfig(profile='collab-aws', region='ap-southeast-2'), + ) + root._fs = fs + + p1, p2 = s3_patches(fs) + with p1, p2, \ + patch.object(get_mod.tmp_mod, 'clone_repo', + return_value=tmp_path / 'clone'), \ + patch.object(get_mod, 'list_source_files', + side_effect=lambda c, p, rev=None: entries[p]), \ + patch.object(hpc, 'require_qxub'), \ + patch.object(hpc, 'submit_workers', return_value=['1.pbs']), \ + patch.object(hpc, 'monitor_jobs', return_value=True): + _, manifest_dir, _ = get_mod.distributed_upload_to_s3( + repository='myrepo', + targets=[('data/fq/AF013-A', 'AF013-A')], + root=root, + num_workers=1, + wait=False, + ) + + raw = (manifest_dir / 'manifest.json').read_text() + metadata = json.loads(raw) + assert metadata['dest_config'] == { + 'profile': 'collab-aws', 'endpoint_url': None, + 'region': 'ap-southeast-2', 'account_id': None, + } + assert 'aws_secret_access_key' not in raw.lower() + assert 'secret' not in raw.lower()