diff --git a/docs/deepvariant-training-case-study.md b/docs/deepvariant-training-case-study.md index 50f9072b..47e5bce8 100644 --- a/docs/deepvariant-training-case-study.md +++ b/docs/deepvariant-training-case-study.md @@ -258,6 +258,11 @@ and also [DirectRunner](https://beam.apache.org/documentation/runners/direct/). Beam can also use other runners, such as [Spark Runner](https://beam.apache.org/documentation/runners/spark/). +NOTE: DirectRunner holds the whole shuffle in memory, so it needs a lot of RAM +for large datasets. If you don't use Google Cloud, see +[tools/aws_shuffle](../tools/aws_shuffle/README.md) for a PySpark alternative +that spills to disk and runs on AWS EMR or a single machine. + First, create a virtual environment to install beam on your machine. ```bash @@ -356,6 +361,11 @@ We can shuffle the validation set locally using `--direct_num_workers=0` sets the number of threads/subprocess to the number of cores of the machine where the pipeline is running. +FACT CHECK: With apache_beam 2.71.0, DirectRunner runs this pipeline on the +Prism runner, which ignores `--direct_num_workers`. On a 64-core machine, +`--direct_num_workers=0` and `--direct_num_workers=1` both took 2:26 for the +validation set and used about 2-3 cores. + ```bash time python3 ${SHUFFLE_SCRIPT_DIR}/shuffle_tfrecords_beam.py \ --project="${YOUR_PROJECT}" \ @@ -365,8 +375,21 @@ time python3 ${SHUFFLE_SCRIPT_DIR}/shuffle_tfrecords_beam.py \ --output_dataset_config_pbtxt="${OUTPUT_DIR}/validation_set.dataset_config.pbtxt" \ --job_name=shuffle-tfrecords \ --runner=DirectRunner \ - --direct_num_workers=0 -``` + --direct_num_workers=0 \ + --job_server_timeout=86400 +``` + +NOTE: Set `--job_server_timeout` to a large number (in seconds) whenever you +use DirectRunner. In recent Beam versions (2.61 and later), DirectRunner runs +the pipeline on Beam's Prism runner, and the Python client gives up on the +job's status stream after `--job_server_timeout` seconds (default 300, raised +from 60 in [apache/beam#36517](https://github.com/apache/beam/pull/36517)). Any +local shuffle that takes longer than 5 minutes then fails with +`StatusCode.DEADLINE_EXCEEDED`, even though the job itself was still running. +The validation set above finishes in about 2.5 minutes, so it never hits the +limit. A local shuffle of the training set takes about 10 minutes and fails at +the 5-minute mark without this flag. Tested with apache_beam 2.71.0, the +version pinned above. Here is the validation_set: diff --git a/tools/aws_shuffle/README.md b/tools/aws_shuffle/README.md new file mode 100644 index 00000000..c0cb96f4 --- /dev/null +++ b/tools/aws_shuffle/README.md @@ -0,0 +1,198 @@ +# Shuffling DeepVariant training examples without Google Cloud + +DeepVariant training needs the examples shuffled globally before training. The +upstream script, `tools/shuffle_tfrecords_beam.py`, supports two ways to run it: + +- **Google Cloud Dataflow.** Scales to large datasets, but needs a GCP account. +- **Beam DirectRunner, on one machine.** Holds the whole shuffle in memory, so + memory grows with the dataset. In our test below, the 1.97 GB (gzipped) + training set needed 32.6 GB of RAM. Larger training sets need a very + large-memory machine or fail with out-of-memory errors. + +`tools/shuffle_tfrecords_spark.py` does the same SHA1-keyed shuffle with plain +PySpark and writes the same outputs (sharded `*.tfrecord.gz`, +`dataset_config.pbtxt`, `example_info.json`). Spark spills the shuffle to disk, +so memory depends on how many cores you use, not on dataset size. It runs on a +single machine or on an AWS EMR cluster. **Plan for disk, not RAM.** + +## Results on real data + +We shuffled the training case study's examples (HG001, made with +`make_examples` from DeepVariant 1.6.1) with both scripts on the same 64-core, +187 GB Linux server: + +- Validation set: 16 files, 0.34 GB gzipped, 59,401 examples. +- Training set: 7 files, 1.97 GB gzipped, 342,758 examples. + +Beam used the v1.10 script with the pinned `apache_beam==2.71.0`. Spark used +PySpark 3.5.3 on Java 17. Peak memory is the total over all processes of the +job. + +**Training set:** + +| | Beam DirectRunner | Spark `local[4]` | Spark `local[*]` | Spark `local[*]`, 64 shards | +|---|---|---|---|---| +| Wall time | 9:52 | 13:09 | 6:33 | **1:45** | +| Peak memory | 32.6 GB | **3.9 GB** | 6.1 GB | 38.9 GB | +| Average cores busy | ~4 | ~4 | ~8 | ~41 | +| Output shards | 5 | 8 | 8 | 64 | + +**Validation set:** + +| | Beam DirectRunner | Spark `local[*]` | +|---|---|---| +| Wall time | 2:26 | 4:08 | +| Peak memory | 5.0-6.0 GB | 3.1 GB | +| Output shards | 2 | 2 | + +Spark at its default settings is slower on small inputs: the default gives 2 +shards for validation and 8 for training, and each shard is sorted and +gzipped by one core (see [Tuning](#tuning)). With `--num_output_shards=64` +the training shuffle took 1:45, with reading done in about 30 seconds. + +### Randomness checks + +We checked every output, Beam and Spark, against its input: + +1. **Same records.** The output holds exactly the input records (compared by + SHA1), and the label counts and `num_examples` in `dataset_config.pbtxt` + match the output. +2. **Order check.** For the first records of each output shard, we printed the + input file and position each record came from. They jump between input + files and positions, with no runs of neighbouring input records. +3. **Neighbour rate.** We measured the share of consecutive output records that + were within 100 records of each other in the same input file, and compared + it with random shuffles of the same records: + + | Output | Neighbour rate | Random shuffles | + |---|---|---| + | Spark, validation | 0.31% | 0.34% | + | Spark, training (8 and 64 shards) | 0.06% | 0.05-0.06% | + | Beam, validation (3 runs) | 0.30-0.36% | 0.29-0.39% (95% range of 200 shuffles) | + | Beam, training | 0.06% | 0.06% | + +All outputs passed all three checks, so both scripts produce a properly random +order. + +### Reproducibility + +The Spark output is deterministic. Each record's shard comes from its SHA1, and +each shard is sorted by SHA1, so the same inputs and shard count always give +the same records in the same order: + +- Two Spark runs on the validation set gave identical records in identical + order in every shard. +- Training runs on 4 cores (`local[4]`) and 64 cores (`local[*]`) gave + identical output in all 8 shards. The core count does not change the result. +- The same shuffle on AWS EMR (2 nodes) and on the local server (64 cores) gave + identical output. +- The `.gz` files themselves differ by checksum between runs, because gzip + stores a timestamp in each file header. Compare decompressed records instead. + +Beam's output order changes on every run, even with identical settings. + +### AWS EMR test + +`run_emr_ec2.sh` shuffled the validation set from S3 on a temporary EMR +cluster: emr-7.5.0, 1 m5.xlarge primary node and 2 m5.2xlarge core nodes with +100 GB gp3 each, using the account's existing EMR roles through +`SERVICE_ROLE` and `INSTANCE_PROFILE`. Provisioning and bootstrap took about +3.5 minutes and the shuffle step about 5.7 minutes, for roughly $0.25. The +output passed the same checks as above, and was identical, record for record +and in the same order, to the output of the same shuffle on the local server. + +## Option 1: one machine + +Any Linux or macOS machine with Java 8, 11 or 17 works, for example a single +EC2 instance. Give it a data volume large enough for the spill (about 1.5x the +gzipped input), and point Spark at that volume with `spark.local.dir`. By +default Spark spills to `/tmp`, which on EC2 is the small root volume. + +```bash +pip install pyspark google-crc32c psutil boto3 +spark-submit --master 'local[*]' --conf spark.local.dir=/mnt/data/spark-tmp \ + tools/shuffle_tfrecords_spark.py \ + --input_pattern_list='s3://my-bucket/out/training_set.with_label.tfrecord-?????-of-?????.gz' \ + --output_pattern_prefix='s3://my-bucket/out/training_set.with_label.shuffled' \ + --output_dataset_name=HG001 \ + --output_dataset_config_pbtxt='s3://my-bucket/out/training_set.dataset_config.pbtxt' +``` + +Inputs and outputs can be `s3://` paths or local paths. If the output goes to +the same disk as the spill, add another 1x the input size of disk space. To cap +memory, use fewer cores, for example `--master 'local[4]'`. + +## Option 2: AWS EMR cluster + +A cluster spreads the same job across several machines, and retries failed +tasks on other nodes, which helps on long runs. `run_emr_ec2.sh` creates a +temporary cluster, runs the shuffle as one Spark step, then terminates the +cluster: + +```bash +BUCKET=s3://my-bucket \ +INPUT='s3://my-bucket/out/training_set.with_label.tfrecord-?????-of-?????.gz' \ +OUT_DIR=s3://my-bucket/out/shuffled NAME=training_set \ +SUBNET_ID=subnet-xxxxxxxx WORKER_COUNT=4 WORKER_EBS_GB=500 \ + tools/aws_shuffle/run_emr_ec2.sh +``` + +The cluster installs the dependencies (`google-crc32c psutil boto3`) through +`emr_bootstrap.sh`. By default it uses the standard EMR roles, +`EMR_DefaultRole` and `EMR_EC2_DefaultRole`, which +`aws emr create-default-roles` creates. Creating them needs IAM permissions, +so ask your AWS admin if you don't have them. If your account uses other EMR +roles, set `SERVICE_ROLE` and `INSTANCE_PROFILE`. To use an existing cluster +that was created with that bootstrap action, set `CLUSTER_ID=j-...`. On EMR, +Spark already spills to the EBS volumes under `/mnt`, so +`WORKER_COUNT x WORKER_EBS_GB` is the total spill space. + +## Resource needs + +- **Memory** depends on the number of tasks running at once (about 0.6 GB each) + plus the Spark JVM, not on dataset size. On the training set, 4 cores peaked + at 3.9 GB and 64 cores with 64 shards at 38.9 GB. +- **Disk** for the spill is about 1.5x the gzipped input. This was measured on + 80,000 synthetic DeepVariant-sized examples (3.0 GB gzipped, about 12 GB + raw), which spilled 4.4 GB. + +## Tuning + +- `--num_output_shards` (default: about 256 MB of gzipped input per shard). + Each shard is sorted and gzipped by one core, so with fewer shards than cores + the last stage leaves cores idle. Set it to at least the number of cores for + speed. The shard count changes which records go into which shard, so keep it + fixed if you need identical output across runs. +- `spark.python.worker.memory` (Spark's default is 512m; the EMR script sets + 1g) is the threshold at which the Python-side sort spills to disk. +- Input records are CRC-checked by default. Use `--no_verify_crc` to skip that. + +## Notes on Beam DirectRunner + +With `apache_beam` 2.61 and later, DirectRunner runs the shuffle pipeline on +Beam's Prism runner. We found two things to watch for: + +- **Long jobs fail** with `StatusCode.DEADLINE_EXCEEDED` once they run past + `--job_server_timeout` (300 s by default in 2.71). Set it to a large value, + for example `--job_server_timeout=86400`. The validation set (2:26) passes + without it; the training set (9:52) failed at 5 minutes. +- **`--direct_num_workers` has no effect.** On the validation set, `=0` (all 64 + cores) and `=1` both took 2:26 and kept about 2-3 cores busy. + +## Why not run the Beam script on Spark? + +Beam's Python `SparkRunner` is a *portable* runner: the Python process starts a +Java job server (downloaded from Maven at runtime), and every Spark executor +launches a Beam SDK harness, by default in a Docker container. On EMR that +breaks in several ways: + +- `spark-submit my_beam_pipeline.py` does not turn a Beam pipeline into a Spark + job; the script just runs as a PySpark driver that then tries to start a + second (Beam) runner. +- The Python `SparkRunner` cannot target `--spark_master=yarn` directly. It only + supports `local[*]` and `spark://host:port`, or building a jar with + `--output_executable_path` and spark-submitting that jar. +- Executors need Docker (or a hand-installed SDK `boot` binary with + `--environment_type=PROCESS`), plus internet access for pip/Maven. + +The PySpark script avoids all of that. diff --git a/tools/aws_shuffle/emr_bootstrap.sh b/tools/aws_shuffle/emr_bootstrap.sh new file mode 100755 index 00000000..d8dabb96 --- /dev/null +++ b/tools/aws_shuffle/emr_bootstrap.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# EMR on EC2 bootstrap action: install the (small) Python deps of +# shuffle_tfrecords_spark.py on every node. No Beam, no TensorFlow needed. +set -euo pipefail +sudo python3 -m pip install --no-cache-dir google-crc32c psutil boto3 diff --git a/tools/aws_shuffle/run_emr_ec2.sh b/tools/aws_shuffle/run_emr_ec2.sh new file mode 100755 index 00000000..07749901 --- /dev/null +++ b/tools/aws_shuffle/run_emr_ec2.sh @@ -0,0 +1,87 @@ +#!/bin/bash +# Shuffle DeepVariant examples on an EMR on EC2 cluster (YARN). +# +# 1) Creates a transient cluster that installs deps via emr_bootstrap.sh, runs +# the shuffle as a single Spark step, then terminates itself. +# 2) Set CLUSTER_ID to submit to an existing cluster instead (it must have been +# created with emr_bootstrap.sh as a bootstrap action). +# +# Disk, not RAM, is what a large shuffle needs: Spark spills the shuffle to the +# EBS volumes of core/task nodes. Give the cluster in total at least 1.5x the +# gzipped size of your input examples (WORKER_COUNT x WORKER_EBS_GB). Output +# goes to S3, so it needs no extra local disk. +set -euo pipefail + +BUCKET="${BUCKET:?set BUCKET, e.g. s3://my-bucket}" +INPUT="${INPUT:-${BUCKET}/training-case-study/output/training_set.with_label.tfrecord-?????-of-?????.gz}" +OUT_DIR="${OUT_DIR:-${BUCKET}/training-case-study/output/shuffled}" +NAME="${NAME:-training_set}" +DATASET_NAME="${DATASET_NAME:-HG001}" +REGION="${REGION:-us-east-1}" +PROFILE="${PROFILE:-}" # optional AWS CLI profile +CLUSTER_ID="${CLUSTER_ID:-}" +# Leave empty to size shards automatically (~256 MB compressed each). +NUM_OUTPUT_SHARDS="${NUM_OUTPUT_SHARDS:-0}" + +# Transient-cluster settings (ignored when CLUSTER_ID is set). +# Default roles come from `aws emr create-default-roles`. +LOG_URI="${LOG_URI:-${BUCKET}/emr-logs/}" +SUBNET_ID="${SUBNET_ID:-}" # required unless CLUSTER_ID is set +INSTANCE_PROFILE="${INSTANCE_PROFILE:-EMR_EC2_DefaultRole}" +SERVICE_ROLE="${SERVICE_ROLE:-EMR_DefaultRole}" +WORKER_TYPE="${WORKER_TYPE:-m5.2xlarge}" +WORKER_COUNT="${WORKER_COUNT:-4}" +WORKER_EBS_GB="${WORKER_EBS_GB:-500}" + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SHUFFLE_PY="${SCRIPT_DIR}/../shuffle_tfrecords_spark.py" +[[ -f "${SHUFFLE_PY}" ]] || SHUFFLE_PY="${SCRIPT_DIR}/shuffle_tfrecords_spark.py" +CODE_PREFIX="${BUCKET}/code/aws_shuffle" +AWS=(aws --region "${REGION}" ${PROFILE:+--profile "${PROFILE}"}) + +"${AWS[@]}" s3 cp "${SHUFFLE_PY}" "${CODE_PREFIX}/" +"${AWS[@]}" s3 cp "${SCRIPT_DIR}/emr_bootstrap.sh" "${CODE_PREFIX}/" + +# JSON (not shorthand) so '?' / '*' / ',' in args survive intact. +STEPS_JSON="$(mktemp)" +cat > "${STEPS_JSON}" <-00000-of-000NN.tfrecord.gz ... + (optional) + /example_info.json (copied if found) + +Python deps on driver and executors: google-crc32c, boto3 (only for s3://). + +Run locally: + spark-submit --master 'local[*]' tools/shuffle_tfrecords_spark.py \ + --input_pattern_list="/tmp/training_set.with_label.tfrecord-?????-of-00016.gz" \ + --output_pattern_prefix="/tmp/out/training_set.with_label.shuffled" \ + --output_dataset_name="HG001" \ + --output_dataset_config_pbtxt="/tmp/out/training_set.dataset_config.pbtxt" + +Run on EMR (see tools/aws_shuffle/README.md for cluster setup): + spark-submit --deploy-mode cluster s3://BUCKET/shuffle_tfrecords_spark.py \ + --input_pattern_list="s3://BUCKET/training_set.with_label.tfrecord-?????-of-?????.gz" \ + --output_pattern_prefix="s3://BUCKET/shuffled/training_set.with_label.shuffled" \ + --output_dataset_name="HG001" \ + --output_dataset_config_pbtxt="s3://BUCKET/shuffled/training_set.dataset_config.pbtxt" +""" +# pylint: enable=line-too-long + +import argparse +import collections +import fnmatch +import gzip +import hashlib +import io +import logging +import os +import re +import shutil +import struct +import textwrap + +COMMENT_HEADER = """# +# --input_pattern_list={} +# --output_pattern_prefix={} +# +""" + +# Default target size of one output shard, measured on compressed input bytes. +_DEFAULT_BYTES_PER_SHARD = 256 * 1024 * 1024 +# S3 multipart part size. Must be >= 5 MiB. +_S3_PART_SIZE = 64 * 1024 * 1024 + + +def parse_cmdline(argv=None): + parser = argparse.ArgumentParser(description=__doc__.split('\n')[0]) + parser.add_argument( + '--input_pattern_list', + required=True, + help='Comma-separated list of TFRecord filename patterns (s3:// or local).') + parser.add_argument( + '--output_pattern_prefix', + required=True, + help='Filename prefix for the output TFRecords.') + parser.add_argument( + '--output_dataset_config_pbtxt', + help='Optional. If set, write a DeepVariantDatasetConfig pbtxt here.') + parser.add_argument( + '--output_dataset_name', + help='Dataset name written into the pbtxt. Required with the pbtxt flag.') + parser.add_argument( + '--num_output_shards', + type=int, + default=0, + help='Number of output shards (= shuffle partitions). 0 means derive it ' + 'from total input size and --target_shard_mb.') + parser.add_argument( + '--target_shard_mb', + type=int, + default=_DEFAULT_BYTES_PER_SHARD // (1024 * 1024), + help='Approximate compressed MB per output shard when ' + '--num_output_shards=0.') + parser.add_argument( + '--no_verify_crc', + action='store_true', + help='Skip CRC verification of input records.') + args = parser.parse_args(argv) + if args.output_dataset_config_pbtxt and not args.output_dataset_name: + parser.error('Need to set --output_dataset_name.') + return args + + +# ---------------------------------------------------------------------------- +# Filesystem helpers: s3:// via boto3, everything else via the local filesystem. +# ---------------------------------------------------------------------------- + + +def _split_s3(path): + m = re.match(r'^s3a?://([^/]+)/?(.*)$', path) + if not m: + raise ValueError('Not an s3 path: {}'.format(path)) + return m.group(1), m.group(2) + + +def _is_s3(path): + return path.startswith('s3://') or path.startswith('s3a://') + + +def _s3_client(): + import boto3 # pylint: disable=g-import-not-at-top + return boto3.client('s3') + + +def glob_with_sizes(pattern): + """Returns a sorted list of (path, size_bytes) matching a glob pattern.""" + if not _is_s3(pattern): + import glob # pylint: disable=g-import-not-at-top + return sorted((p, os.path.getsize(p)) for p in glob.glob(pattern)) + bucket, key_pattern = _split_s3(pattern) + # List from the longest literal prefix, then filter with fnmatch. + literal_prefix = re.split(r'[*?\[]', key_pattern, maxsplit=1)[0] + paginator = _s3_client().get_paginator('list_objects_v2') + out = [] + for page in paginator.paginate(Bucket=bucket, Prefix=literal_prefix): + for obj in page.get('Contents', []): + if fnmatch.fnmatchcase(obj['Key'], key_pattern): + out.append(('s3://{}/{}'.format(bucket, obj['Key']), obj['Size'])) + return sorted(out) + + +def open_read(path): + if _is_s3(path): + bucket, key = _split_s3(path) + return _s3_client().get_object(Bucket=bucket, Key=key)['Body'] + return open(path, 'rb') + + +def copy_file(src, dst): + if _is_s3(src) and _is_s3(dst): + sb, sk = _split_s3(src) + db, dk = _split_s3(dst) + _s3_client().copy({'Bucket': sb, 'Key': sk}, db, dk) + elif not _is_s3(src) and not _is_s3(dst): + os.makedirs(os.path.dirname(dst) or '.', exist_ok=True) + shutil.copyfile(src, dst) + else: + with open_read(src) as f: + write_bytes(dst, f.read()) + + +def write_bytes(path, data): + if _is_s3(path): + bucket, key = _split_s3(path) + _s3_client().put_object(Bucket=bucket, Key=key, Body=data) + else: + os.makedirs(os.path.dirname(path) or '.', exist_ok=True) + with open(path, 'wb') as f: + f.write(data) + + +def exists(path): + if _is_s3(path): + bucket, key = _split_s3(path) + resp = _s3_client().list_objects_v2(Bucket=bucket, Prefix=key, MaxKeys=1) + return any(o['Key'] == key for o in resp.get('Contents', [])) + return os.path.exists(path) + + +class _S3MultipartWriter(io.RawIOBase): + """Write-only file object that streams to S3 with multipart upload.""" + + def __init__(self, path): + super().__init__() + self._bucket, self._key = _split_s3(path) + self._client = _s3_client() + self._buf = bytearray() + self._parts = [] + self._upload_id = self._client.create_multipart_upload( + Bucket=self._bucket, Key=self._key)['UploadId'] + + def writable(self): + return True + + def write(self, b): + self._buf.extend(b) + if len(self._buf) >= _S3_PART_SIZE: + self._flush_part() + return len(b) + + def _flush_part(self): + n = len(self._parts) + 1 + resp = self._client.upload_part( + Bucket=self._bucket, Key=self._key, UploadId=self._upload_id, + PartNumber=n, Body=bytes(self._buf)) + self._parts.append({'PartNumber': n, 'ETag': resp['ETag']}) + self._buf = bytearray() + + def close(self): + if self.closed: + return + try: + if self._buf or not self._parts: + self._flush_part() + self._client.complete_multipart_upload( + Bucket=self._bucket, Key=self._key, UploadId=self._upload_id, + MultipartUpload={'Parts': self._parts}) + except Exception: + self.abort() + raise + super().close() + + def abort(self): + self._client.abort_multipart_upload( + Bucket=self._bucket, Key=self._key, UploadId=self._upload_id) + super().close() + + +def open_write(path): + if _is_s3(path): + return _S3MultipartWriter(path) + os.makedirs(os.path.dirname(path) or '.', exist_ok=True) + return open(path, 'wb') + + +# ---------------------------------------------------------------------------- +# TFRecord framing. Each record is: +# uint64 length | uint32 masked_crc32c(length) | data | uint32 masked_crc32c(data) +# ---------------------------------------------------------------------------- + + +def _masked_crc(data): + import google_crc32c # pylint: disable=g-import-not-at-top + crc = google_crc32c.value(data) + return (((crc >> 15) | (crc << 17)) + 0xa282ead8) & 0xffffffff + + +def read_tfrecords(path, verify_crc=True): + """Yields raw record bytes from one (optionally gzipped) TFRecord file.""" + raw = open_read(path) + f = gzip.GzipFile(fileobj=raw, mode='rb') if path.endswith('.gz') else raw + try: + while True: + header = f.read(12) + if not header: + return + if len(header) != 12: + raise IOError('Truncated record header in {}'.format(path)) + length, length_crc = struct.unpack('> 3, tag & 7 + if wire == 0: + value, pos = _read_varint(buf, pos) + elif wire == 1: + value, pos = buf[pos:pos + 8], pos + 8 + elif wire == 2: + n, pos = _read_varint(buf, pos) + value, pos = buf[pos:pos + n], pos + n + elif wire == 5: + value, pos = buf[pos:pos + 4], pos + 4 + else: + raise ValueError('Unsupported protobuf wire type {}'.format(wire)) + yield field, wire, value + + +def example_label(record): + """Returns features['label'].int64_list.value[0], or None if absent.""" + buf = memoryview(record) + for f1, _, features in _iter_fields(buf): # Example.features = 1 + if f1 != 1: + continue + for f2, _, entry in _iter_fields(features): # Features.feature = 1 (map) + if f2 != 1: + continue + key = feature = None + for f3, _, v in _iter_fields(entry): + if f3 == 1: + key = bytes(v) + elif f3 == 2: + feature = v + if key != b'label' or feature is None: + continue + for f4, _, int64_list in _iter_fields(feature): # Feature.int64_list = 3 + if f4 != 3: + continue + for f5, wire, v in _iter_fields(int64_list): # Int64List.value = 1 + if f5 != 1: + continue + if wire == 2: # packed + v, _ = _read_varint(v, 0) + return v - (1 << 64) if v >= (1 << 63) else v + return None + + +# ---------------------------------------------------------------------------- +# Pipeline. +# ---------------------------------------------------------------------------- + + +def _shard_path(prefix, index, num_shards): + return '{}-{:05d}-of-{:05d}.tfrecord.gz'.format(prefix, index, num_shards) + + +def _partition_of(key): + return int.from_bytes(key[:8], 'big') + + +def make_config_string(name, tfrecord_path, num_examples): + return textwrap.dedent(""" + name: "{}" + tfrecord_path: "{}-?????-of-?????.tfrecord.gz" + num_examples: {} + """.format(name, tfrecord_path, num_examples)) + + +def copy_example_info_json(input_pattern, output_pattern_prefix): + matches = glob_with_sizes('{}*example_info.json'.format(input_pattern)) + if not matches: + return + dst = os.path.join( + os.path.dirname(output_pattern_prefix), 'example_info.json') + if not exists(dst): + logging.info('Copying %s to %s', matches[0][0], dst) + copy_file(matches[0][0], dst) + + +def main(argv=None): + args = parse_cmdline(argv) + from pyspark.sql import SparkSession # pylint: disable=g-import-not-at-top + + spark = SparkSession.builder.appName('shuffle-tfrecords').getOrCreate() + sc = spark.sparkContext + + patterns = [p.strip() for p in args.input_pattern_list.split(',') if p.strip()] + copy_example_info_json(patterns[0], args.output_pattern_prefix) + + inputs = [] + for pattern in patterns: + matched = glob_with_sizes(pattern) + if not matched: + raise ValueError('No files match {}'.format(pattern)) + inputs.extend(matched) + total_bytes = sum(size for _, size in inputs) + num_shards = args.num_output_shards or max( + 1, -(-total_bytes // (args.target_shard_mb * 1024 * 1024))) + logging.info('Shuffling %d input files (%.2f GB compressed) into %d shards.', + len(inputs), total_bytes / 1e9, num_shards) + + verify_crc = not args.no_verify_crc + output_prefix = args.output_pattern_prefix + + def write_shard(index, sorted_pairs): + labels = collections.Counter() + writer = open_write(_shard_path(output_prefix, index, num_shards)) + try: + with gzip.GzipFile(fileobj=writer, mode='wb') as gz: + for _, record in sorted_pairs: + write_tfrecord(gz, record) + labels[example_label(record)] += 1 + except BaseException: + if hasattr(writer, 'abort'): + writer.abort() + else: + writer.close() + raise + writer.close() + yield labels + + paths = [p for p, _ in inputs] + label_counts = ( + sc.parallelize(paths, len(paths)) + .flatMap(lambda p: read_tfrecords(p, verify_crc)) + .map(lambda r: (hashlib.sha1(r).digest(), r)) + .repartitionAndSortWithinPartitions( + numPartitions=num_shards, partitionFunc=_partition_of) + .mapPartitionsWithIndex(write_shard) + .reduce(lambda a, b: a + b)) + num_examples = sum(label_counts.values()) + logging.info('Wrote %d examples. Per label: %s', num_examples, + dict(label_counts)) + + if args.output_dataset_config_pbtxt: + per_label = ''.join( + '# class{}: {}\n'.format(k, v) + for k, v in sorted(label_counts.items(), key=lambda kv: str(kv[0]))) + summary = ('# Generated by shuffle_tfrecords_spark.py\n' + + COMMENT_HEADER.format(args.input_pattern_list, output_prefix) + + per_label + + make_config_string(args.output_dataset_name, output_prefix, + num_examples)) + write_bytes(args.output_dataset_config_pbtxt, summary.encode('utf-8')) + spark.stop() + + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO) + main()