From 66f41162d6c368e15c1e9b6562756c8bea6ebb5c Mon Sep 17 00:00:00 2001 From: John Reeves Date: Fri, 21 Aug 2026 00:07:18 +1000 Subject: [PATCH] feat(get): --workers streams to S3 across compute nodes (0.25.0) `dt get` had no distributed mode. `dt push --workers` had solved multi-node transfer already, and the machinery was never push-specific, so this reuses it rather than rewriting it (#172). Scoped to `s3://` destinations, which is a decision rather than an unfinished edge: -o s3://... bound by the network -> a second node adds bandwidth -o some/dir/ bound by the shared FS -> a second node adds contention A local destination on Lustre is limited by the filesystem, not the node, so `--workers` is a hard usage error there pointing at `-j`. On NCI there is a second reason to like the S3 case: `copyq` is the queue with outbound network access, and it is already hpc.DEFAULT_TRANSFER_QUEUE. Reuse, per the issue: - `hpc.partition_by_size()` is the LPT bin-packer lifted out of `push.partition_manifest`, which was pure bin-packing with nothing push- specific in it. What stays behind in push is the part that genuinely is: a file is named by its hash, so its size has to be read from the local cache. The docstring says where the boundary is and why LPT belongs to qxub workers and not to a ThreadPoolExecutor -- separate processes cannot steal work, a shared queue self-corrects, and freezing an assignment the queue would fix is strictly worse. - `hpc.save_manifest()` takes a `metadata` dict. Everything a worker needs beyond its own partition travels that way, because it is a fresh process on another machine: push needs a remote name, get also needs a destination, credentials config and the flags the run was invoked with. - `submit_workers` / `monitor_jobs` / `load_worker_partition` unchanged. The design questions #172 raised, settled: *Reachability.* Resolve once on the submitting node and pass paths. The submitter has to resolve every row anyway -- LPT needs byte sizes and `list_source_files` is where sizes come from -- so having paid for that it ships the resolved md5, size and destination key per file. **No worker ever runs `dvc list`**, which is what keeps N compute nodes off the shared clone's SQLite state db; `_run_dvc_list`'s retry loop exists because that lock is real. A worker opens the clone once, read-only, for one thing: the source remote's filesystem and credentials out of its `.dvc/config`. Link types and the destination identity are likewise resolved on the submitter 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. *Partition unit.* Per file, not per row -- the issue framed this as a trade against per-row reporting, but it is not one. Each task carries its row, so reporting stays per-row and in CSV order even when a sample directory spans nodes, and per-file balances properly in the two cases per-row cannot: a single row of thousands of files, and rows differing by an order of magnitude, which fastq samples reliably do. *Reporting.* Workers write `result_.json`; the submitter collates by row tag. A row whose files do not all come back is failed even when every result that did arrive succeeded -- a job killed on walltime would otherwise read as a clean run of a short row. Each worker preflights for itself, so `--dest-account-id` is enforced on the node that actually writes. The manifest carries the profile *name*, endpoint and region; no key material. Also, from the issue's note about the ThreadPoolExecutor idiom being copy-pasted five times: this adds a sixth use, so per "factor it then" the two in `get.py` are now behind `_place_all` and `_upload_all`, and `_resolve_rows` is shared by the local and S3 row-resolve paths. That makes `upload_to_s3` resolve rows concurrently, which it did not before. A cross-module `utils.run_parallel()` is still not worth it. Not verified here, and it needs a real submission: whether more nodes actually move more bytes. #172 asks for that measurement and it remains open -- if the *source* remote is on /g/data the read side is a shared filesystem again and the same contention argument applies to it. docs/get.md says so and says to time a single-node `-j 8` run first. Tested: 30 new unit tests. The central one replaces `hpc.submit_workers` with a stub that runs the worker in-process, so the manifest is exercised as a real interface rather than asserted against its spelling. Co-Authored-By: Claude --- docs/commands.md | 3 +- docs/config_options.md | 5 +- docs/get.md | 67 ++++ dt/__init__.py | 2 +- dt/cli.py | 143 ++++++- dt/get.py | 479 ++++++++++++++++++++--- dt/hpc.py | 122 +++++- dt/push.py | 44 +-- pyproject.toml | 2 +- tests/unit/test_get_workers.py | 671 +++++++++++++++++++++++++++++++++ 10 files changed, 1431 insertions(+), 107 deletions(-) create mode 100644 tests/unit/test_get_workers.py 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()