-
Notifications
You must be signed in to change notification settings - Fork 1
upload-to-s3: read the source once instead of three times #77
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,6 +3,12 @@ set -euo pipefail | |
|
|
||
| bin="$(dirname "$0")" | ||
|
|
||
| # Temp artifacts, cleaned up however the script exits. Globals so the EXIT trap | ||
| # sees them regardless of where in main() we leave (return, `exit`, or set -e). | ||
| tmp="" | ||
| workdir="" | ||
| trap '[[ -n "$workdir" ]] && rm -rf "$workdir"; [[ -n "$tmp" ]] && rm -f "$tmp"' EXIT | ||
|
|
||
| main() { | ||
| local quiet=0 | ||
|
|
||
|
|
@@ -20,28 +26,62 @@ main() { | |
| local dst="${2:?A destination s3:// URL is required as the second argument.}" | ||
| local cloudfront_domain="${3:-}" | ||
|
|
||
| # Fail early if $src is unreadable. Otherwise `tee < "$src"` below fails, but | ||
| # the background hasher/counter -- already blocked opening their fifos -- get | ||
| # no writer and hang instead of being cleaned up. | ||
| [[ -r "$src" ]] || { echo "upload-to-s3: cannot read source file: $src" >&2; exit 1; } | ||
|
|
||
| local s3path="${dst#s3://}" | ||
| local bucket="${s3path%%/*}" | ||
| local key="${s3path#*/}" | ||
|
|
||
| local src_hash dst_hash no_hash=0000000000000000000000000000000000000000000000000000000000000000 | ||
| src_hash="$("$bin/sha256sum" < "$src")" | ||
| dst_hash="$(aws s3api head-object --bucket "$bucket" --key "$key" --query Metadata.sha256sum --output text 2>/dev/null || echo "$no_hash")" | ||
| # Read $src exactly once: while compressing it to a temp file, tee the same | ||
| # bytes to the hasher and the line counter. The old code read $src three | ||
| # separate times (sha256sum, wc -l, zstd); for the large aligned/sequences | ||
| # files those reads -- not the S3 transfer -- dominated the upload rules. | ||
| # | ||
| # Uploading a real file (rather than piping to `aws s3 cp -`) also lets the | ||
| # CLI parallelize the multipart transfer; a stdin stream can't seek and | ||
| # degrades to one slow stream. The temp lives next to $src (same big volume). | ||
| local upload_file="$src" | ||
| workdir="$(mktemp -d)" | ||
|
|
||
| if [[ $src_hash != "$dst_hash" ]]; then | ||
| # The record count may have changed | ||
| local -a compressor=() | ||
| case "$dst" in | ||
| *.gz) compressor=(gzip -c) ;; | ||
| *.xz) compressor=(xz -2 -T0 -c) ;; | ||
| *.zst) compressor=(zstd -T0 -c) ;; | ||
| esac | ||
|
|
||
| local src_hash src_record_count | ||
| if [[ ${#compressor[@]} -gt 0 ]]; then | ||
| tmp="$(mktemp "$src.upload.XXXXXX")" | ||
| upload_file="$tmp" | ||
| mkfifo "$workdir/hash.fifo" "$workdir/count.fifo" | ||
| "$bin/sha256sum" < "$workdir/hash.fifo" > "$workdir/hash" & | ||
| local hash_pid=$! | ||
| wc -l < "$workdir/count.fifo" > "$workdir/count" & | ||
| local count_pid=$! | ||
| tee "$workdir/hash.fifo" "$workdir/count.fifo" < "$src" | "${compressor[@]}" > "$tmp" | ||
| # Wait one at a time: `wait a b` returns only b's status, so a failing | ||
| # hasher would be masked from set -e (uploading with a bad sha256sum). | ||
| wait "$hash_pid" | ||
| wait "$count_pid" | ||
|
Comment on lines
+68
to
+69
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. With this spinning out 2 processes in the background, do we have to worry about oversubscribing CPUs when Snakemake schedules multiple upload jobs? I think we're fine based on the improved runtimes documented in the PR description, but something to keep in mind in the future.
Member
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. My (partial) understanding here:
|
||
| src_hash="$(cat "$workdir/hash")" | ||
| src_record_count="$(cat "$workdir/count")" | ||
| else | ||
| # No compression (e.g. the small *.json files) -- hash and count directly. | ||
| src_hash="$("$bin/sha256sum" < "$src")" | ||
| src_record_count="$(wc -l < "$src")" | ||
| fi | ||
|
|
||
| local dst_hash no_hash=0000000000000000000000000000000000000000000000000000000000000000 | ||
| dst_hash="$(aws s3api head-object --bucket "$bucket" --key "$key" --query Metadata.sha256sum --output text 2>/dev/null || echo "$no_hash")" | ||
|
|
||
| if [[ $src_hash != "$dst_hash" ]]; then | ||
| echo "Uploading $src → $dst" | ||
| if [[ "$dst" == *.gz ]]; then | ||
| gzip -c "$src" | ||
| elif [[ "$dst" == *.xz ]]; then | ||
| xz -2 -T0 -c "$src" | ||
| elif [[ "$dst" == *.zst ]]; then | ||
| zstd -T0 -c "$src" | ||
| else | ||
| cat "$src" | ||
| fi | aws s3 cp --no-progress - "$dst" --metadata sha256sum="$src_hash",recordcount="$src_record_count" "$(content-type "$dst")" | ||
|
|
||
| aws s3 cp --no-progress "$upload_file" "$dst" --metadata sha256sum="$src_hash",recordcount="$src_record_count" "$(content-type "$dst")" | ||
|
|
||
| if [[ -n $cloudfront_domain ]]; then | ||
| echo "Creating CloudFront invalidation for $cloudfront_domain/$key" | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.