From 7abdb4d28c392b0f28c0b73a649103d01de00376 Mon Sep 17 00:00:00 2001 From: Muzi Li Date: Wed, 15 Jul 2026 23:03:51 +0000 Subject: [PATCH 1/2] feat(data): add Bagz format support with GCS FUSE compatibility Implements end-to-end support for the Bagz (.bagz) dataset format inside the MaxText training and data processing pipeline. Enables highly efficient, POSIX-compliant file loading optimized specifically for GCS FUSE and local filesystem environments. Scope of Changes: - Input Pipeline: Integrated upstream `grain.BagzDataSource` for direct, multiprocess-safe reading of Bagz file shards via PyGrain. - Data Processing: Registered `bagz` as a first-class supported file type across PyGrain data loaders, feature normalizers, and tokenizers. - Tooling: Created `download_hf_dataset_as_bagz.py` with multi-worker support, HuggingFace streaming, and resilient checkpointing for GCSFUSE mounts. Testing: - Unit Tests: Added comprehensive unit test suite in `tests/unit/bagz_data_processing_test.py` covering RandomAccessDataSource compatibility, MapDataset, ElasticIterator (single and multi-worker), and end-to-end `get_datasets` integration. - End-to-End Local Smoke Test: Executed a 3-step offline training loop on a local CPU environment using the newly generated Bagz dataset via `DECOUPLE_GCLOUD=TRUE`. - Parallel Data Generation: Successfully converted and streamed HF datasets into multi-shard Bagz files (38 shards, Salesforce/wikitext). --- .../base_requirements/requirements.txt | 1 + .../requirements/requirements.txt | 1 + src/maxtext/configs/base.yml | 2 +- src/maxtext/configs/types.py | 6 +- .../input_pipeline/data_processing_utils.py | 2 +- .../input_pipeline/grain_data_processing.py | 9 +- .../input_pipeline/input_pipeline_utils.py | 2 + .../trainers/tokenizer/train_tokenizer.py | 9 +- tests/unit/bagz_data_processing_test.py | 124 ++++++++ .../download_hf_dataset_as_bagz.py | 287 ++++++++++++++++++ 10 files changed, 434 insertions(+), 9 deletions(-) create mode 100644 tests/unit/bagz_data_processing_test.py create mode 100644 tools/data_generation/download_hf_dataset_as_bagz.py diff --git a/src/dependencies/requirements/base_requirements/requirements.txt b/src/dependencies/requirements/base_requirements/requirements.txt index fe6e05daf3..012ebdceb5 100644 --- a/src/dependencies/requirements/base_requirements/requirements.txt +++ b/src/dependencies/requirements/base_requirements/requirements.txt @@ -12,6 +12,7 @@ google-cloud-aiplatform google-cloud-mldiagnostics google-cloud-monitoring grain[parquet] +bagz huggingface_hub>=1.14.0 jax jaxlib diff --git a/src/dependencies/requirements/requirements.txt b/src/dependencies/requirements/requirements.txt index 05c2be074b..7112daf5c0 100644 --- a/src/dependencies/requirements/requirements.txt +++ b/src/dependencies/requirements/requirements.txt @@ -12,6 +12,7 @@ google-cloud-aiplatform google-cloud-mldiagnostics>=0.5.10 google-cloud-monitoring grain[parquet] +bagz huggingface_hub jax!=0.7.1 jaxlib!=0.7.1 diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 6b23e31d40..fe37eb1e78 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -763,7 +763,7 @@ hf_access_token: '' grain_train_files: '' grain_eval_files: '' grain_train_mixture_config_path: '' # Path to a JSON file specifying the mixture weights for Grain training data. -grain_file_type: 'arrayrecord' # arrayrecord or parquet +grain_file_type: 'arrayrecord' # arrayrecord, parquet, tfrecord, or bagz grain_packing_type: 'first_fit' # 'first_fit', 'best_fit' or 'concat_then_split'. See details of the corresponding module in https://google-grain.readthedocs.io/en/latest/grain.experimental.html grain_worker_count: 1 # Set to -1 to enable auto-tuning: automatically determines optimal worker count. See https://google-grain.readthedocs.io/en/latest/_autosummary/grain.experimental.pick_performance_config.html grain_per_worker_buffer_size: 1 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 72cc3de96a..d1d50e72b7 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1409,7 +1409,7 @@ class GrainDataset(BaseModel): ) grain_file_type: str = Field( "arrayrecord", - description="File type for Grain data. Supported: arrayrecord, tfrecord, parquet.", + description="File type for Grain data. Supported: arrayrecord, tfrecord, parquet, bagz.", ) grain_use_elastic_iterator: bool = Field( False, @@ -3386,9 +3386,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "Colocated python data input is only supported with Pathways (single" " controller) enabled (`enable_single_controller=True`)." ) - if self.grain_use_elastic_iterator and self.grain_file_type != "arrayrecord": + if self.grain_use_elastic_iterator and self.grain_file_type not in ("arrayrecord", "bagz"): raise ValueError( - "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord`. " + "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord` or `bagz`. " "tfrecord and parquet pipelines use `InterleaveIterDataset` (a many-to-one " "IterDataset transform), which `ElasticIterator` forbids. " f"Got grain_file_type={self.grain_file_type}." diff --git a/src/maxtext/input_pipeline/data_processing_utils.py b/src/maxtext/input_pipeline/data_processing_utils.py index 3336b1fed3..d46f12629c 100644 --- a/src/maxtext/input_pipeline/data_processing_utils.py +++ b/src/maxtext/input_pipeline/data_processing_utils.py @@ -27,7 +27,7 @@ def parse_and_keep_features(dataset, config, data_columns, tokenize): """Parse arrayrecord features or keep specified columns for other formats.""" - if config.grain_file_type in ("arrayrecord", "tfrecord"): + if config.grain_file_type in ("arrayrecord", "tfrecord", "bagz"): dataset = dataset.map(input_pipeline_utils.ParseFeatures(data_columns, tokenize)) dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(data_columns, tokenize)) else: diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 220c3ce82d..8c4f648330 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/src/maxtext/input_pipeline/grain_data_processing.py @@ -98,11 +98,14 @@ def get_datasets( elastic=False, ): """Load dataset from array_record files for using with grain""" - if data_file_type == "arrayrecord": + if data_file_type in ("arrayrecord", "bagz"): # Helper function to find files, create data source, and wrap in MapDataset def create_dataset_from_pattern(pattern): files = find_data_files(pattern) - source = grain.ArrayRecordDataSource(files) + if data_file_type == "arrayrecord": + source = grain.ArrayRecordDataSource(files) + elif data_file_type == "bagz": + source = grain.BagzDataSource(files) return grain.MapDataset.source(source) # Handle mixture config with named datasets, allows flexibility in recovering checkpoints @@ -215,7 +218,7 @@ def create_dataset_from_pattern(pattern): return dataset else: raise ValueError( - f"grain pipeline supports (arrayrecord, tfrecord, parquet) as grain_file_type, but got {data_file_type}" + f"grain pipeline supports (arrayrecord, tfrecord, parquet, bagz) as grain_file_type, but got {data_file_type}" ) diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index d996a88908..b40740c7bb 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -1070,3 +1070,5 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]: element[f"{self.data_column}_mrope_deltas"] = mrope_position_deltas return element + + diff --git a/src/maxtext/trainers/tokenizer/train_tokenizer.py b/src/maxtext/trainers/tokenizer/train_tokenizer.py index fdd25fef61..69817c2327 100644 --- a/src/maxtext/trainers/tokenizer/train_tokenizer.py +++ b/src/maxtext/trainers/tokenizer/train_tokenizer.py @@ -41,6 +41,7 @@ import jax import grain.python as grain +import bagz from maxtext.input_pipeline import input_pipeline_utils from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT @@ -93,6 +94,12 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys: dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True)) dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True)) return iter(dataset) + elif data_file_type == "bagz": + source = input_pipeline_utils.BagzDataSource(data_files) + dataset = grain.MapDataset.source(source) + dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True)) + dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True)) + return iter(dataset) elif data_file_type == "tfrecord": dataset = grain.MapDataset.source(data_files) dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset) @@ -103,7 +110,7 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys: dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True)) return iter(dataset) else: - raise ValueError(f"Unsupported grain_file_type: {data_file_type!r}. Use 'parquet', 'arrayrecord', or 'tfrecord'.") + raise ValueError(f"Unsupported grain_file_type: {data_file_type!r}. Use 'parquet', 'arrayrecord', 'tfrecord', or 'bagz'.") def _dump_chars_to_textfile(dataset_iter: Iterator, maxchars: int = int(1e7), data_keys=("text",)) -> tuple[str, int]: diff --git a/tests/unit/bagz_data_processing_test.py b/tests/unit/bagz_data_processing_test.py new file mode 100644 index 0000000000..d72762f922 --- /dev/null +++ b/tests/unit/bagz_data_processing_test.py @@ -0,0 +1,124 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for Bagz data processing and ElasticIterator compatibility.""" + +import os +import sys +import tempfile +import unittest +from absl import flags +try: + import pytest + cpu_only = pytest.mark.cpu_only +except ImportError: + cpu_only = lambda x: x +import bagz +import grain.python as grain +from grain.experimental import ElasticIterator + +from grain.python import BagzDataSource + +flags.FLAGS(sys.argv) + + +@cpu_only +class GrainBagzProcessingTest(unittest.TestCase): + """Test BagzDataSource and compatibility with Grain MapDataset and ElasticIterator.""" + + def setUp(self): + super().setUp() + self.test_dir = tempfile.TemporaryDirectory() + self.bagz_path = os.path.join(self.test_dir.name, "test_data.bagz") + self.num_records = 20 + + writer = bagz.Writer(self.bagz_path) + for i in range(self.num_records): + writer.write(f"record_{i}".encode("utf-8")) + writer.close() + + def tearDown(self): + self.test_dir.cleanup() + super().tearDown() + + def test_bagz_data_source_random_access(self): + source = BagzDataSource([self.bagz_path]) + self.assertEqual(len(source), self.num_records) + self.assertEqual(source[0].decode("utf-8"), "record_0") + self.assertEqual(source[self.num_records - 1].decode("utf-8"), f"record_{self.num_records - 1}") + + def test_bagz_map_dataset(self): + source = BagzDataSource([self.bagz_path]) + ds = grain.MapDataset.source(source).map(lambda x: x.decode("utf-8")) + self.assertEqual(len(ds), self.num_records) + self.assertEqual(ds[5], "record_5") + + def test_bagz_elastic_iterator_single_process(self): + source = BagzDataSource([self.bagz_path]) + ds = grain.MapDataset.source(source).map(lambda x: x.decode("utf-8")) + + iter_ds = ElasticIterator( + ds, + global_batch_size=4, + shard_options=grain.ShardOptions(shard_index=0, shard_count=1), + ) + it = iter(iter_ds) + batch1 = next(it) + batch2 = next(it) + self.assertEqual(len(batch1), 4) + self.assertEqual(len(batch2), 4) + self.assertEqual(list(batch1), ["record_0", "record_1", "record_2", "record_3"]) + self.assertEqual(list(batch2), ["record_4", "record_5", "record_6", "record_7"]) + + def test_bagz_elastic_iterator_multi_process(self): + source = BagzDataSource([self.bagz_path]) + ds = grain.MapDataset.source(source).map(lambda x: x.decode("utf-8")) + + mp_options = grain.MultiprocessingOptions(num_workers=2, per_worker_buffer_size=1) + iter_ds = ElasticIterator( + ds, + global_batch_size=4, + shard_options=grain.ShardOptions(shard_index=0, shard_count=1), + multiprocessing_options=mp_options, + ) + it = iter(iter_ds) + batch1 = next(it) + batch2 = next(it) + self.assertEqual(len(batch1), 4) + self.assertEqual(len(batch2), 4) + + def test_bagz_get_datasets_integration(self): + from maxtext.input_pipeline import grain_data_processing + train_ds = grain_data_processing.get_datasets( + data_file_pattern=self.bagz_path, + data_file_type="bagz", + shuffle=False, + shuffle_seed=0, + shuffle_buffer_size=1, + num_epoch=1, + dataloading_host_index=0, + dataloading_host_count=1, + grain_worker_count=0, + grain_num_threads=1, + grain_prefetch_buffer_size=1, + grain_data_source_max_workers=1, + elastic=False, + ) + records = list(train_ds) + self.assertEqual(len(records), self.num_records) + self.assertEqual(records[0], b"record_0") + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/data_generation/download_hf_dataset_as_bagz.py b/tools/data_generation/download_hf_dataset_as_bagz.py new file mode 100644 index 0000000000..049ca67bd3 --- /dev/null +++ b/tools/data_generation/download_hf_dataset_as_bagz.py @@ -0,0 +1,287 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Download a HuggingFace dataset via streaming and save as Bagz files. +Only supports text dataset for now. + +Supports writing to local path, including GCS bucket mounted via GCSFUSE. +Produces filenames: {name}-XXXXX-of-YYYYY.bagz +User can control file-size-mb and workers for parallelism. + +GCS output recommendation: + Mount the bucket with write-optimized GCSFUSE flags and pass the mount path as --output. +""" + +import argparse +import json +import multiprocessing +import os +import pathlib +import shutil +import sys +import time + +from maxtext.input_pipeline.protos import example_pb2 +from maxtext.input_pipeline.protos import feature_pb2 + +from datasets import load_dataset +import requests +import bagz +MAX_RETRIES = 10 +RETRY_WAIT_SECONDS = 30 + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Convert a HuggingFace streaming dataset to Bagz files.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument("--dataset", required=True, help="HuggingFace dataset name (e.g. Salesforce/wikitext)") + parser.add_argument("--config", default=None, help="Dataset config as a JSON string for load_dataset") + parser.add_argument("--split", default="train", help="Dataset split to convert (default: train)") + parser.add_argument("--output", required=True, help="Output directory (local path or GCSFUSE mount).") + parser.add_argument("--name-prefix", default=None, help="Filename prefix (default: derived from dataset)") + parser.add_argument("--file-size-mb", type=float, default=1000.0, help="Target file size per Bagz shard in MB") + parser.add_argument("--workers", type=int, default=None, help="Number of concurrent worker processes") + parser.add_argument("--token", default=None, help="HuggingFace auth token for gated datasets") + return parser.parse_args() + + +def _to_feature(value): + if isinstance(value, bool): + return feature_pb2.Feature(int64_list=feature_pb2.Int64List(value=[int(value)])) + if isinstance(value, int): + return feature_pb2.Feature(int64_list=feature_pb2.Int64List(value=[value])) + if isinstance(value, float): + return feature_pb2.Feature(float_list=feature_pb2.FloatList(value=[value])) + if isinstance(value, (str, bytes)): + v = value.encode("utf-8") if isinstance(value, str) else value + return feature_pb2.Feature(bytes_list=feature_pb2.BytesList(value=[v])) + if isinstance(value, list): + if not value: + return feature_pb2.Feature(bytes_list=feature_pb2.BytesList(value=[])) + first = value[0] + if isinstance(first, (bool, int)): + return feature_pb2.Feature(int64_list=feature_pb2.Int64List(value=[int(v) for v in value])) + if isinstance(first, float): + return feature_pb2.Feature(float_list=feature_pb2.FloatList(value=value)) + if isinstance(first, (str, bytes)): + v = [x.encode("utf-8") if isinstance(x, str) else x for x in value] + return feature_pb2.Feature(bytes_list=feature_pb2.BytesList(value=v)) + + json_v = json.dumps(value, ensure_ascii=False).encode("utf-8") if value is not None else b"" + return feature_pb2.Feature(bytes_list=feature_pb2.BytesList(value=[json_v] if json_v else [])) + + +def serialize_example(example: dict) -> bytes: + features = {k: _to_feature(v) for k, v in example.items()} + return example_pb2.Example(features=feature_pb2.Features(feature=features)).SerializeToString() + + +def process_shard(task): + worker_id = task["worker_id"] + num_workers = task["num_workers"] + dataset_name = task["dataset"] + config = task["config"] + split = task["split"] + token = task["token"] + output_dir = task["output"] + file_size_bytes = task["file_size_bytes"] + checkpoint_path = os.path.join(output_dir, f".checkpoint-worker-{worker_id:04d}.json") + + kwargs = {"path": dataset_name, "streaming": True, "split": split, "token": token} + if config: + kwargs.update(json.loads(config)) + + def load_and_restore(): + ds = load_dataset(**kwargs).shard(num_shards=num_workers, index=worker_id) + if os.path.exists(checkpoint_path): + with open(checkpoint_path, "r", encoding="utf-8") as f: + ckpt = json.load(f) + ds.load_state_dict(ckpt["ds_state"]) + return ds, ckpt + return ds, { + "local_file_idx": 0, + "total_bytes": 0, + "total_written": 0, + "filenames": [], + "adjusted_file_size_bytes": file_size_bytes, + } + + for attempt in range(MAX_RETRIES): + try: + ds, state = load_and_restore() + break + except (requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout) as e: + print(f"Worker {worker_id}: Connection failed (attempt {attempt+1}/{MAX_RETRIES}): {e}") + if attempt + 1 == MAX_RETRIES: + raise + time.sleep(RETRY_WAIT_SECONDS) + + local_file_idx, total_bytes, total_written, filenames, adjusted_file_size_bytes = ( + state[k] for k in ["local_file_idx", "total_bytes", "total_written", "filenames", "adjusted_file_size_bytes"] + ) + + file_bytes, writer, t0 = 0, None, time.time() + + def open_writer(): + nonlocal writer + fname = f"worker-{worker_id:04d}-{local_file_idx:06d}.bagz" + writer = bagz.Writer(os.path.join(output_dir, fname)) + filenames.append(fname) + + def close_writer(): + nonlocal writer + if writer: + writer.close() + writer = None + + def save_checkpoint(): + ckpt = { + "worker_id": worker_id, + "local_file_idx": local_file_idx, + "total_bytes": total_bytes, + "total_written": total_written, + "filenames": filenames, + "adjusted_file_size_bytes": adjusted_file_size_bytes, + "ds_state": ds.state_dict(), + } + with open(checkpoint_path, "w", encoding="utf-8") as f: + json.dump(ckpt, f) + + open_writer() + for attempt in range(MAX_RETRIES): + try: + for example in ds: + record = serialize_example(example) + record_size = len(record) + writer.write(record) + file_bytes += record_size + total_bytes += record_size + total_written += 1 + + if file_bytes >= adjusted_file_size_bytes: + close_writer() + actual_size = os.path.getsize(os.path.join(output_dir, filenames[-1])) + if actual_size > 0: + adjusted_file_size_bytes = int(file_size_bytes / (actual_size / file_bytes)) + + elapsed = time.time() - t0 + speed = (file_bytes / 1024 / 1024) / elapsed + print(f"Worker {worker_id}: Finished shard {local_file_idx} ({actual_size/1024/1024:.1f} MB) in {elapsed:.1f}s ({speed:.1f} MB/s)") + + local_file_idx += 1 + file_bytes = 0 + save_checkpoint() + open_writer() + t0 = time.time() + + close_writer() + if file_bytes == 0: + filenames.pop() + else: + actual_size = os.path.getsize(os.path.join(output_dir, filenames[-1])) + elapsed = time.time() - t0 + speed = (file_bytes / 1024 / 1024) / elapsed + print(f"Worker {worker_id}: Finished FINAL shard {local_file_idx} ({actual_size/1024/1024:.1f} MB) in {elapsed:.1f}s ({speed:.1f} MB/s)") + save_checkpoint() + + if os.path.exists(checkpoint_path): + os.remove(checkpoint_path) + return filenames + + except (requests.exceptions.ConnectionError, requests.exceptions.ReadTimeout) as e: + print(f"Worker {worker_id}: Connection lost during iteration (attempt {attempt+1}/{MAX_RETRIES}): {e}") + close_writer() + if attempt + 1 == MAX_RETRIES: + raise + time.sleep(RETRY_WAIT_SECONDS) + + return filenames + + +def rename_files(all_filenames, output_dir, prefix): + print("\nVerifying downloaded files...") + existing_files = [] + for fname in all_filenames: + fpath = os.path.join(output_dir, fname) + if os.path.exists(fpath) and os.path.getsize(fpath) > 0: + existing_files.append(fname) + else: + print(f" [Warning] File not found or empty, skipping: {fpath}") + + total_files = len(existing_files) + if total_files == 0: + print("No valid files to rename.") + return + + print(f"\nRenaming {total_files} verified files...") + existing_files.sort() + width = max(5, len(str(total_files))) + + for idx, old_name in enumerate(existing_files): + new_name = f"{prefix}-{idx:0{width}d}-of-{total_files:0{width}d}.bagz" + if old_name != new_name: + os.rename(os.path.join(output_dir, old_name), os.path.join(output_dir, new_name)) + + print(f" Renamed to: {prefix}-00000-of-{total_files:0{width}d}.bagz ... {prefix}-{total_files - 1:0{width}d}-of-{total_files:0{width}d}.bagz") + + +def convert(args): + if args.output.startswith("gs://"): + raise ValueError("gs:// paths are not supported. Mount the bucket with GCSFUSE and pass the mount path instead.") + num_workers = args.workers or os.cpu_count() + file_size_bytes = int(args.file_size_mb * 1024 * 1024) + + dataset_name_safe = args.dataset.replace("/", "_") + prefix = args.name_prefix or dataset_name_safe + output_dir = os.path.abspath(args.output) + os.makedirs(output_dir, exist_ok=True) + + tasks = [] + for i in range(num_workers): + tasks.append({ + "worker_id": i, + "num_workers": num_workers, + "dataset": args.dataset, + "config": args.config, + "split": args.split, + "token": args.token, + "output": output_dir, + "file_size_bytes": file_size_bytes, + }) + + start_time = time.time() + print(f"Starting conversion of '{args.dataset}' with {num_workers} workers...") + + if num_workers == 1: + results = [process_shard(tasks[0])] + else: + with multiprocessing.Pool(processes=num_workers) as pool: + results = pool.map(process_shard, tasks) + + all_filenames = [] + for worker_filenames in results: + all_filenames.extend(worker_filenames) + + rename_files(all_filenames, output_dir, prefix) + + total_time = time.time() - start_time + print(f"\nSuccessfully converted dataset in {total_time/60:.1f} minutes.") + + +if __name__ == "__main__": + convert(parse_args()) From cbd0bc271d1ee4961af2c0f6d40583a3b8174594 Mon Sep 17 00:00:00 2001 From: Muzi Li Date: Tue, 4 Aug 2026 22:23:11 +0000 Subject: [PATCH 2/2] docs: add example usage commands to download_hf_dataset_as_bagz docstring Add practical example commands to the docstring of `download_hf_dataset_as_bagz.py`: - Basic local directory conversion (Salesforce/wikitext) - Write-optimized GCSFUSE mounted bucket output with custom shard prefix and file size - Private/gated HuggingFace dataset download using auth token --- .../download_hf_dataset_as_bagz.py | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/tools/data_generation/download_hf_dataset_as_bagz.py b/tools/data_generation/download_hf_dataset_as_bagz.py index 049ca67bd3..c6275285a0 100644 --- a/tools/data_generation/download_hf_dataset_as_bagz.py +++ b/tools/data_generation/download_hf_dataset_as_bagz.py @@ -12,14 +12,40 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Download a HuggingFace dataset via streaming and save as Bagz files. +"""Download a HuggingFace dataset via streaming and save as Bagz files. Only supports text dataset for now. Supports writing to local path, including GCS bucket mounted via GCSFUSE. Produces filenames: {name}-XXXXX-of-YYYYY.bagz User can control file-size-mb and workers for parallelism. +Examples: + 1. Basic conversion of a HuggingFace dataset to local directory: + python3 tools/data_generation/download_hf_dataset_as_bagz.py \ + --dataset=Salesforce/wikitext \ + --config='{"name": "wikitext-103-raw-v1"}' \ + --split=train \ + --output=/tmp/wikitext_bagz \ + --file-size-mb=100.0 \ + --workers=4 + + 2. Writing to a write-optimized GCSFUSE mounted bucket with custom prefix: + python3 tools/data_generation/download_hf_dataset_as_bagz.py \ + --dataset=allenai/c4 \ + --config='{"name": "en"}' \ + --split=train \ + --output=/gcs_mount/my_bucket/c4_bagz \ + --name-prefix=c4_en_train \ + --file-size-mb=1000.0 \ + --workers=16 + + 3. Converting a gated/private HuggingFace dataset using an auth token: + python3 tools/data_generation/download_hf_dataset_as_bagz.py \ + --dataset=meta-llama/Llama-2-7b \ + --split=train \ + --output=/tmp/llama_data \ + --token=$HF_TOKEN + GCS output recommendation: Mount the bucket with write-optimized GCSFUSE flags and pass the mount path as --output. """