From 90f0b27c856c4d738691ba885bfb7490df65ce58 Mon Sep 17 00:00:00 2001 From: yinshiyi Date: Tue, 13 Oct 2026 13:54:06 -0700 Subject: [PATCH 1/6] Add PySpark TFRecord shuffler for AWS EMR and single machines shuffle_tfrecords_beam.py runs on Google Cloud Dataflow or on Beam's DirectRunner. DirectRunner holds the whole shuffle in memory, so large training sets need a very large-memory machine, and Beam's Python Spark runner is impractical on EMR (no direct YARN submission, Docker needed on workers, runtime downloads). tools/shuffle_tfrecords_spark.py is a plain PySpark version of the same SHA1-keyed shuffle with the same outputs (sharded tfrecord.gz, dataset_config.pbtxt, example_info.json). Inputs are streamed and Spark's sort-based shuffle spills to disk, so memory does not grow with the dataset. tools/aws_shuffle/ adds a launch script for a transient EMR cluster. --- docs/deepvariant-training-case-study.md | 5 + tools/aws_shuffle/README.md | 101 ++++++ tools/aws_shuffle/emr_bootstrap.sh | 5 + tools/aws_shuffle/run_emr_ec2.sh | 86 +++++ tools/shuffle_tfrecords_spark.py | 449 ++++++++++++++++++++++++ 5 files changed, 646 insertions(+) create mode 100644 tools/aws_shuffle/README.md create mode 100755 tools/aws_shuffle/emr_bootstrap.sh create mode 100755 tools/aws_shuffle/run_emr_ec2.sh create mode 100644 tools/shuffle_tfrecords_spark.py diff --git a/docs/deepvariant-training-case-study.md b/docs/deepvariant-training-case-study.md index 50f9072b..b2fce7d9 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 diff --git a/tools/aws_shuffle/README.md b/tools/aws_shuffle/README.md new file mode 100644 index 00000000..ed5d125c --- /dev/null +++ b/tools/aws_shuffle/README.md @@ -0,0 +1,101 @@ +# Shuffling DeepVariant training examples without Google Cloud + +The existing `tools/shuffle_tfrecords_beam.py` runs on Google Cloud Dataflow or +on Beam's DirectRunner. Without Google Cloud, that leaves DirectRunner, which +holds the whole shuffle in memory: a large training set needs a machine with +very large RAM, or fails with out-of-memory errors. + +`tools/shuffle_tfrecords_spark.py` does the same SHA1-keyed shuffle with plain +PySpark and produces the same outputs (sharded `*.tfrecord.gz`, +`dataset_config.pbtxt`, `example_info.json`). Spark spills the shuffle to disk, +so memory stays roughly constant no matter how big the dataset is. **Plan for +disk, not RAM:** about 1.5x the gzipped input size (see measurements below). + +There are two ways to run it: + +- **One machine** (recommended). Simplest; a large dataset just takes longer. +- **An AWS EMR cluster**, if you don't want to wait on a single instance. + +## 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. + +## Option 2: AWS EMR cluster (faster) + +A cluster spreads the same job across several machines, so it finishes sooner. +It also 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`. It uses the standard EMR roles (create them once with +`aws emr create-default-roles`). 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 (measured) + +Test: 80,000 DeepVariant-sized examples, 3.0 GB gzipped (~12 GB raw), +`local[4]` with `--driver-memory 1g` and `spark.python.worker.memory=256m`. + +| Peak | Value | +|---|---| +| Total memory, all Spark processes | 2.1 GB | +| Largest Python worker | 0.46 GB | +| Spill disk | 4.4 GB (~1.5x gzipped input) | +| Open files per Python worker | 37 | + +Memory does not grow with dataset size. Disk does: plan for ~1.5x the gzipped +input for spill, plus 1x more if the output is written to the same disk. + +## Tuning + +- `--num_output_shards` (default: about 256 MB of compressed input per shard). + Each shard is one Spark partition, so you need enough shards to keep every + core busy. +- `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. + +## 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..6948657c --- /dev/null +++ b/tools/aws_shuffle/run_emr_ec2.sh @@ -0,0 +1,86 @@ +#!/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 roughly 3x the +# compressed size of your input examples. +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() From 4ad5c09eb4130439f5dff45150ba9b49e704ff12 Mon Sep 17 00:00:00 2001 From: yinshiyi Date: Tue, 13 Oct 2026 15:21:20 -0700 Subject: [PATCH 2/6] Document Beam DirectRunner 5-minute job_server_timeout failure With apache_beam 2.61+, DirectRunner runs pipelines on the Prism runner. The Python client applies --job_server_timeout as a deadline on the whole job-status stream, so any local shuffle that runs longer than the timeout fails with StatusCode.DEADLINE_EXCEEDED while the job is still running. The default is 300 s; it was raised from 60 s in https://github.com/apache/beam/pull/36517, whose description recommends setting --job_server_timeout for longer jobs. The case study's validation-set shuffle finishes in about 2.5 minutes, so it never hits the limit. Shuffling the training set locally takes about 10 minutes: it failed at 5:03 with the default, and completed in 9:52 with --job_server_timeout=86400 (apache_beam 2.71.0, r1.10 script). Add --job_server_timeout=86400 to the DirectRunner example and a note recommending a large value for any local Beam run. --- docs/deepvariant-training-case-study.md | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/docs/deepvariant-training-case-study.md b/docs/deepvariant-training-case-study.md index b2fce7d9..22620c27 100644 --- a/docs/deepvariant-training-case-study.md +++ b/docs/deepvariant-training-case-study.md @@ -370,8 +370,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: From 5b2ce95c08db4cef5949b9e39be9d4361555f8d7 Mon Sep 17 00:00:00 2001 From: yinshiyi Date: Tue, 13 Oct 2026 15:34:42 -0700 Subject: [PATCH 3/6] Fact-check --direct_num_workers=0 claim for Prism-based DirectRunner The case study says --direct_num_workers=0 sets the number of threads/subprocesses to the number of cores. With apache_beam 2.71.0, DirectRunner runs the shuffle pipeline on the Prism runner, and only the older FnApiRunner reads this option. A/B test on the validation set (64-core machine, same inputs): --direct_num_workers=1: 2:26.1 wall, 232% CPU, PrismRunner --direct_num_workers=0: 2:26.8 wall, 232% CPU, PrismRunner --- docs/deepvariant-training-case-study.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/deepvariant-training-case-study.md b/docs/deepvariant-training-case-study.md index 22620c27..47e5bce8 100644 --- a/docs/deepvariant-training-case-study.md +++ b/docs/deepvariant-training-case-study.md @@ -361,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}" \ From 19cbf0505eef16e5c52c739819e60483ba01dbd0 Mon Sep 17 00:00:00 2001 From: yinshiyi Date: Tue, 13 Oct 2026 16:44:16 -0700 Subject: [PATCH 4/6] Update aws_shuffle README with real-data benchmark and checks Fold the background on why a non-Beam shuffle is needed into the README, and replace the synthetic-only measurements with results from the training case study's examples (HG001; 0.34 GB validation, 1.97 GB training), shuffled with both scripts on the same 64-core server: - Training: Beam DirectRunner 9:52 / 32.6 GB; Spark local[4] 13:09 / 3.9 GB; local[*] 6:33 / 6.1 GB; local[*] with 64 shards 1:45 / 38.9 GB. - All outputs keep exactly the input records and label counts, and their neighbour rate matches random shuffles of the same records. - Spark output is deterministic: repeated runs, and 4- vs 64-core runs, give identical records in identical order. - One validation run on AWS EMR (emr-7.5.0) for about $0.25. Also document the shard-count bottleneck in Tuning and the two Beam DirectRunner pitfalls found (job_server_timeout, direct_num_workers). --- tools/aws_shuffle/README.md | 154 ++++++++++++++++++++++++++++-------- 1 file changed, 122 insertions(+), 32 deletions(-) diff --git a/tools/aws_shuffle/README.md b/tools/aws_shuffle/README.md index ed5d125c..5117d4e9 100644 --- a/tools/aws_shuffle/README.md +++ b/tools/aws_shuffle/README.md @@ -1,20 +1,101 @@ # Shuffling DeepVariant training examples without Google Cloud -The existing `tools/shuffle_tfrecords_beam.py` runs on Google Cloud Dataflow or -on Beam's DirectRunner. Without Google Cloud, that leaves DirectRunner, which -holds the whole shuffle in memory: a large training set needs a machine with -very large RAM, or fails with out-of-memory errors. +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 produces the same outputs (sharded `*.tfrecord.gz`, +PySpark and writes the same outputs (sharded `*.tfrecord.gz`, `dataset_config.pbtxt`, `example_info.json`). Spark spills the shuffle to disk, -so memory stays roughly constant no matter how big the dataset is. **Plan for -disk, not RAM:** about 1.5x the gzipped input size (see measurements below). +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: -There are two ways to run it: +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: -- **One machine** (recommended). Simplest; a large dataset just takes longer. -- **An AWS EMR cluster**, if you don't want to wait on a single instance. + | 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 `.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 + +The same script 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. Provisioning and bootstrap took about 4 minutes and the +shuffle step 5.5 minutes, for roughly $0.25. The output had the same 59,401 +records and label counts as the input. ## Option 1: one machine @@ -34,14 +115,15 @@ spark-submit --master 'local[*]' --conf spark.local.dir=/mnt/data/spark-tmp \ ``` 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. +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 (faster) +## Option 2: AWS EMR cluster -A cluster spreads the same job across several machines, so it finishes sooner. -It also 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: +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 \ @@ -58,30 +140,38 @@ 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 (measured) - -Test: 80,000 DeepVariant-sized examples, 3.0 GB gzipped (~12 GB raw), -`local[4]` with `--driver-memory 1g` and `spark.python.worker.memory=256m`. - -| Peak | Value | -|---|---| -| Total memory, all Spark processes | 2.1 GB | -| Largest Python worker | 0.46 GB | -| Spill disk | 4.4 GB (~1.5x gzipped input) | -| Open files per Python worker | 37 | +## Resource needs -Memory does not grow with dataset size. Disk does: plan for ~1.5x the gzipped -input for spill, plus 1x more if the output is written to the same disk. +- **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 compressed input per shard). - Each shard is one Spark partition, so you need enough shards to keep every - core busy. +- `--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 From 26893aa6a9f98eedf7cd06bbe3daeff7543063a1 Mon Sep 17 00:00:00 2001 From: yinshiyi Date: Tue, 13 Oct 2026 16:54:10 -0700 Subject: [PATCH 5/6] Align EMR disk-sizing comment with the README The run_emr_ec2.sh header asked for about 3x the compressed input in disk, while the README measures the spill at about 1.5x the gzipped input (plus 1x only when output lands on the same disk). On EMR the output goes to S3, so state 1.5x and point at WORKER_COUNT x WORKER_EBS_GB. --- tools/aws_shuffle/run_emr_ec2.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tools/aws_shuffle/run_emr_ec2.sh b/tools/aws_shuffle/run_emr_ec2.sh index 6948657c..07749901 100755 --- a/tools/aws_shuffle/run_emr_ec2.sh +++ b/tools/aws_shuffle/run_emr_ec2.sh @@ -7,8 +7,9 @@ # 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 roughly 3x the -# compressed size of your input examples. +# 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}" From 6155e45baf028907a1abd7029d4c0b5c39d7972b Mon Sep 17 00:00:00 2001 From: yinshiyi Date: Tue, 13 Oct 2026 17:12:07 -0700 Subject: [PATCH 6/6] Document EMR role overrides and the committed launcher's EMR test Ran the committed run_emr_ec2.sh end-to-end on the validation set (emr-7.5.0, 2 x m5.2xlarge core nodes), overriding SERVICE_ROLE and INSTANCE_PROFILE with the account's existing EMR roles because the standard EMR_DefaultRole / EMR_EC2_DefaultRole did not exist there. Provisioning took about 3.5 minutes and the step about 5.7 minutes. The output kept all 59,401 records and label counts, and matched the local 64-core Spark run record for record and in order. README changes: - Option 2: name the default roles, note that creating them needs IAM permissions, and point to SERVICE_ROLE / INSTANCE_PROFILE overrides. - AWS EMR test: describe this run of the committed launcher. - Reproducibility: EMR and local outputs are identical. --- tools/aws_shuffle/README.md | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/tools/aws_shuffle/README.md b/tools/aws_shuffle/README.md index 5117d4e9..c0cb96f4 100644 --- a/tools/aws_shuffle/README.md +++ b/tools/aws_shuffle/README.md @@ -84,6 +84,8 @@ the same records in the same order: 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. @@ -91,11 +93,13 @@ Beam's output order changes on every run, even with identical settings. ### AWS EMR test -The same script shuffled the validation set from S3 on a temporary EMR +`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. Provisioning and bootstrap took about 4 minutes and the -shuffle step 5.5 minutes, for roughly $0.25. The output had the same 59,401 -records and label counts as the input. +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 @@ -134,11 +138,14 @@ SUBNET_ID=subnet-xxxxxxxx WORKER_COUNT=4 WORKER_EBS_GB=500 \ ``` The cluster installs the dependencies (`google-crc32c psutil boto3`) through -`emr_bootstrap.sh`. It uses the standard EMR roles (create them once with -`aws emr create-default-roles`). 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. +`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