Add BagzDataSource for BagZ file format support

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/

[Marlon666]

Overview

This PR introduces BagzDataSource, a first-class RandomAccessDataSource in PyGrain for reading BagZ (.bagz) datasets.
While native C++ bagz.Reader can be used directly in simple single-process scripts, integrating BagZ into distributed multiprocessing training pipelines requires additional handling for pickling, fork-safety, and path list normalization. BagzDataSource provides an official, zero-overhead bridge to address these requirements.

Key Features & Motivation

1. Multiprocessing / Fork-Safety (Lazy Initialization):
   In distributed multi-worker pipelines (grain_worker_count > 0 or ElasticIterator), PyGrain serializes (pickle) the data source across worker processes. Because open C++ bagz.Reader handles cannot be pickled (TypeError: cannot pickle 'bagz.lib.bagz.Reader' object) and sharing open file descriptors across forked workers is unsafe on POSIX / GCS FUSE mounts, BagzDataSource implements custom __getstate__/__setstate__ to reset reader handles and lazily instantiating bagz.Reader on demand inside each worker.
2. Input Path Type Normalization:
   Supports passing a single path string/FileInstruction as well as a list/sequence of paths (transparently normalizing sequences into comma-separated strings required by bagz.Reader).
3. Native Grain Observability & Tracing:
   Integrates with Grain's API initialization counter (_api_usage_counter) and read-stage monitoring (@dataset_stats.trace_input_pipeline).

Testing

• Added comprehensive unit tests in grain/_src/python/data_sources_test.py (BagzDataSourceTest), covering:
  • Subclass compliance with RandomAccessDataSource
  • Empty sequence validation
  • Accurate random-access indexing
  • Full pickle serialization/deserialization verification across instantiated reader states
• All 13 tests in data_sources_test.py pass cleanly.

---

📚 Documentation preview 📚: https://google-grain--1376.org.readthedocs.build/