Custom dataloader

.coveragerc

@@ -1,5 +1,72 @@
+[run]
+# Source directories to measure coverage for
+source = src
+
+# Parallel mode for multiprocessing support
+parallel = True
+
+# Branch coverage measurement
+branch = True
+
+# Data file location
+data_file = .coverage
+
 [report]
-include = ./src/megatron/energon/**
+# Include source code in the report
+include = src/megatron/energon/**
+
+# Minimum coverage percentage to pass
+fail_under = 80
+
+# Show missing lines in the report
+show_missing = True
+
+# Skip covered files in the report
+skip_covered = False
+
+# Skip empty files
+skip_empty = True
+
+# Precision for coverage percentages
+precision = 2
+
+# Sort order for the report
+sort = filename
+
+# Exclude lines from coverage
+exclude_lines =
+    # Have to re-enable the standard pragma
+    pragma: no cover
+
+    # Don't complain about missing debug-only code:
+    def __repr__
+    if self\.debug
+
+    # Don't complain if tests don't hit exceptions:
+    raise .*
+
+    # Don't complain if non-runnable code isn't run:
+    if 0:
+    if __name__ == .__main__.:
+
+    # Don't complain about abstract methods, they aren't run:
+    @(abc\.)?abstractmethod
+
+    # Don't complain about type checking code:
+    if TYPE_CHECKING:
+
+
+[html]
+# Directory to put the HTML report
+directory = htmlcov
+
+# Title for the HTML report
+title = Megatron Energon Coverage Report
 
 [xml]
-output = ./coverage.xml
+# Output file for XML report
+output = coverage.xml
+
+[lcov]
+# Output file for LCOV report
+output = lcov.info

justfile

@@ -25,7 +25,15 @@ check: dev-sync
 
 # Execute all unit tests
 test: dev-sync
-    uv run -m unittest discover -v -s tests
+    uv run pytest tests -v
+
+coverage: dev-sync
+    COVERAGE_PROCESS_START=.coveragerc uv run -m coverage run --parallel-mode --concurrency=multiprocessing -m pytest tests
+    # COVERAGE_PROCESS_START=.coveragerc uv run -m coverage run --parallel-mode --concurrency=multiprocessing -m pytest tests/test_dataloader.py
+    # COVERAGE_PROCESS_START=.coveragerc uv run -m coverage run --parallel-mode --concurrency=multiprocessing -m pytest tests/test_typed_converter.py
+    uv run -m coverage combine
+    uv run -m coverage lcov
+    echo "Coverage LCOV report generated at ./lcov.info"
 
 # Build the docs
 docs: dev-sync

pyproject.toml

@@ -37,6 +37,8 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = [
+  "coverage",
+  "pytest",
   "ruff",
   "sphinxcontrib-napoleon",
   "sphinx",
@@ -114,3 +116,17 @@ exclude = [
   "docs",
 ]
 
+[tool.coverage.run]
+branch = true
+parallel = true
+concurrency = ["multiprocessing"]
+
+[tool.coverage.report]
+show_missing = true
+skip_covered = true
+
+[tool.coverage.html]
+show_contexts = true
+
+[tool.pytest.ini_options]
+addopts = "--tb=native --capture=no"

src/megatron/energon/__init__.py

@@ -70,7 +70,6 @@
     load_dataset,
     prepare_metadataset,
 )
-from megatron.energon.savable_loader import SavableDataLoader
 from megatron.energon.source_info import SourceInfo
 from megatron.energon.task_encoder import (
     AugmentTaskEncoder,
@@ -189,7 +188,6 @@
     "reraise_exception",
     "Sample",
     "SampleDecoder",
-    "SavableDataLoader",
     "SavableDataset",
     "ShuffleBufferDataset",
     "SimilarityInterleavedSample",

src/megatron/energon/cache/base.py

@@ -11,7 +11,7 @@
 T = TypeVar("T")
 
 
-class FileStore(Generic[T]):
+class FileStore(ABC, Generic[T]):
     """Base type for a dataset that can be accessed randomly by sample key."""
 
     @abstractmethod
@@ -40,6 +40,21 @@ def get_path(self) -> str:
         """Returns the path to the dataset."""
         ...
 
+    @abstractmethod
+    def worker_init(self) -> None:
+        """Initializes the file store for the current worker."""
+        raise NotImplementedError("worker_init is not implemented for this file store")
+
+    @abstractmethod
+    def worker_close(self) -> None:
+        """Closes the file store for the current worker."""
+        raise NotImplementedError("worker_close is not implemented for this file store")
+
+    @abstractmethod
+    def close(self) -> None:
+        """Closes the file store."""
+        raise NotImplementedError("close is not implemented for this file store")
+
     def get_media_metadata(self, key: str) -> MediaMetadataBase:
         """Return the media metadata for the given key if available."""
 
@@ -73,6 +88,15 @@ def _decode_raw(self, data: T, **kwargs) -> T:
         """
         return self._inner._decode_raw(data, **kwargs)
 
+    def worker_init(self) -> None:
+        self._inner.worker_init()
+
+    def worker_close(self) -> None:
+        self._inner.worker_close()
+
+    def close(self) -> None:
+        self._inner.close()
+
 
 @edataclass
 class Lazy(Generic[T]):
@@ -169,6 +193,20 @@ def get_lazy(self, ds: FileStore, fname: str) -> Lazy:
         """
         ...
 
+    @abstractmethod
+    def worker_init(self) -> None:
+        """
+        Initialize the cache pool for the current worker.
+        """
+        ...
+
+    @abstractmethod
+    def worker_close(self) -> None:
+        """
+        Close the cache pool for the current worker.
+        """
+        ...
+
     @abstractmethod
     def to_cache(self, data: T, name: str) -> Lazy[T]:
         """

src/megatron/energon/cache/file_cache_pool.py

@@ -157,6 +157,8 @@ class FileStoreCachePool(CachePool, ForkMixin):
     # Whether the pool is shutting down
     _shutting_down: bool = False
 
+    _workers_initialized: dict[int, bool] = {}
+
     def __init__(
         self,
         *,
@@ -278,6 +280,9 @@ def _cache_out_task(self, ds: FileStore, fname: str, entry: _PendingTask) -> boo
         with self._lock:
             if self._shutting_down:
                 return False
+        if not self._workers_initialized.get(threading.get_ident(), False):
+            ds.worker_init()
+            self._workers_initialized[threading.get_ident()] = True
 
         # Perform the data read
         if self.method == "raw":
@@ -363,6 +368,12 @@ def get_lazy(self, ds: FileStore, fname: str) -> FileCacheLazy:
 
         return FileCacheLazy(ds=ds, fname=fname, pool=self, entry=entry)
 
+    def worker_init(self) -> None:
+        pass
+
+    def worker_close(self) -> None:
+        pass
+
     def to_cache(self, data: T, name: str) -> CacheFileLazy[T]:
         """
         Move the data to the cache and return a lazy to fetch it later.

src/megatron/energon/cache/file_store.py

@@ -83,6 +83,15 @@ def __getitem__(self, key: str) -> tuple[bytes, SourceInfo]:
             file_names=(key,),
         )
 
+    def worker_init(self) -> None:
+        pass
+
+    def worker_close(self) -> None:
+        pass
+
+    def close(self) -> None:
+        pass
+
     def get_path(self) -> str:
         """Returns the path to the dataset."""
         return str(self.base_dir)
@@ -157,7 +166,7 @@ def get_path(self) -> str:
     def get_media_metadata(self, key: str) -> MediaMetadataBase:
         if self._media_metadata_available is None:
             try:
-                has_metadata = self.sqlite_reader.db_has_media_metadata()
+                has_metadata = self._sqlite_reader.db_has_media_metadata()
             except sqlite3.Error as exc:  # pragma: no cover - defensive
                 raise RuntimeError(
                     "Failed to inspect media metadata table. Re-run `energon prepare --media-metadata-by-...`."
@@ -172,7 +181,7 @@ def get_media_metadata(self, key: str) -> MediaMetadataBase:
             self._media_metadata_available = True
 
         try:
-            row = self.sqlite_reader.get_media_metadata(key)
+            row = self._sqlite_reader.get_media_metadata(key)
         except sqlite3.Error as exc:  # pragma: no cover - defensive
             raise RuntimeError(
                 "Failed to load media metadata. Re-run `energon prepare --media-metadata-by-...`."

src/megatron/energon/cache/no_cache.py

@@ -47,6 +47,12 @@ def get(self, ds: FileStore, fname: str, sample: Any = None) -> Any:
     def get_lazy(self, ds: FileStore, fname: str) -> DirectLazy:
         return DirectLazy(ds=ds, fname=fname, pool=self)
 
+    def worker_init(self) -> None:
+        pass
+
+    def worker_close(self) -> None:
+        pass
+
     def to_cache(self, data: T, name: str) -> DirectLazy:
         return MockLazy(fname=name, get_fn=lambda _: data)
 

src/megatron/energon/dataloader/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: BSD-3-Clause
+from .dataloader import DataLoader
+from .pin_memory import NoPinMemory, PinMemory, PinMemoryThread
+from .workers import DataLoaderWorker, ForkDataLoaderWorker, ThreadDataLoaderWorker
+
+__all__ = [
+    "DataLoader",
+    "PinMemory",
+    "NoPinMemory",
+    "PinMemoryThread",
+    "DataLoaderWorker",
+    "ThreadDataLoaderWorker",
+    "ForkDataLoaderWorker",
+]

src/megatron/energon/dataloader/asynchronous/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: BSD-3-Clause
+from .base import Asynchronous, QueueProtocol, WorkerCommand, WorkerResult
+from .fork import ForkAsynchronous
+from .thread import ThreadAsynchronous
+
+__all__ = [
+    "Asynchronous",
+    "QueueProtocol",
+    "WorkerCommand",
+    "WorkerResult",
+    "ForkAsynchronous",
+    "ThreadAsynchronous",
+]