[bugfix] correct dataloader checkpoint resume across forked workers and epochs

tzrec/datasets/dataset.py

@@ -756,6 +756,7 @@ def create_dataloader(
     mode: Mode = Mode.TRAIN,
     gl_cluster: Optional[Dict[str, Union[int, str]]] = None,
     debug_level: int = 0,
+    checkpoint_state: Optional[Dict[str, Any]] = None,
 ) -> DataLoader:
     """Build dataloader.
 
@@ -768,6 +769,11 @@ def create_dataloader(
         gl_cluster (dict, bool): if set, reuse the graphlearn cluster.
         debug_level (int): dataset debug level, when mode=predict and
             debug_level > 0, will dump fg encoded data to debug_str
+        checkpoint_state (dict, optional): dataloader checkpoint state for
+            resume. Must be set here rather than on the returned dataloader:
+            the eager ``iter(dataloader)`` below forks persistent workers,
+            which keep a fork-time copy of the dataset, so state applied
+            afterwards never reaches them.
 
     Return:
         dataloader (dataloader): a DataLoader.
@@ -783,6 +789,8 @@ def create_dataloader(
         mode=mode,
         debug_level=debug_level,
     )
+    if checkpoint_state:
+        dataset.load_state_dict(checkpoint_state)
 
     kwargs = {}
     if data_config.num_workers < 1:

tzrec/datasets/parquet_dataset_test.py

@@ -23,6 +23,7 @@
 from torch import distributed as dist
 from torch.utils.data import DataLoader
 
+from tzrec.datasets.dataset import create_dataloader
 from tzrec.datasets.parquet_dataset import ParquetDataset, ParquetReader, ParquetWriter
 from tzrec.features.feature import create_features
 from tzrec.protos import data_pb2, feature_pb2
@@ -235,6 +236,50 @@ def test_parquet_dataset_checkpoint_resume(self):
                     # New offset should be less than or equal to acc checkpoint offset
                     self.assertLessEqual(new_offset, checkpoint_state_acc[key])
 
+    def test_create_dataloader_checkpoint_state_reaches_workers(self):
+        """State passed to create_dataloader must reach forked workers.
+
+        create_dataloader eagerly starts persistent workers, so state applied
+        to the returned dataloader's dataset afterwards never reaches them;
+        the checkpoint_state argument applies it before the fork.
+        """
+        feature_cfgs = self._create_feature_cfgs()
+        features = create_features(feature_cfgs)
+
+        with tempfile.TemporaryDirectory(prefix="tzrec_") as test_dir:
+            # 20000 rows at max_rows_per_file=5000 -> 4 files, so
+            # data_config.num_workers=2 survives the num_files clamp.
+            self._create_test_parquet_data(test_dir, num_rows=20000)
+            input_path = f"{test_dir}/*.parquet"
+            data_config = data_pb2.DataConfig(
+                batch_size=128,
+                dataset_type=data_pb2.DatasetType.ParquetDataset,
+                fg_mode=data_pb2.FgMode.FG_NONE,
+                label_fields=["label"],
+                num_workers=2,
+            )
+
+            dataloader1 = create_dataloader(data_config, features, input_path)
+            iterator1 = iter(dataloader1)
+            checkpoint_state = {}
+            for _ in range(4):
+                batch1 = next(iterator1)
+                update_dataloder_state(checkpoint_state, batch1.checkpoint_info)
+            self.assertGreater(len(checkpoint_state), 0)
+            del iterator1, dataloader1
+
+            dataloader2 = create_dataloader(
+                data_config, features, input_path, checkpoint_state=checkpoint_state
+            )
+            iterator2 = iter(dataloader2)
+            for _ in range(4):
+                batch2 = next(iterator2)
+                for key, new_offset in batch2.checkpoint_info.items():
+                    if key in checkpoint_state:
+                        # without the pre-fork state, workers replay from row 0
+                        self.assertGreater(new_offset, checkpoint_state[key])
+            del iterator2, dataloader2
+
 
 class ParquetReaderTest(unittest.TestCase):
     def setUp(self):

tzrec/main.py

@@ -588,22 +588,6 @@ def train_and_evaluate(
     # Build feature
     features = _create_features(list(pipeline_config.feature_configs), data_config)
 
-    # Build dataloader
-    train_dataloader = create_dataloader(
-        data_config, features, pipeline_config.train_input_path, mode=Mode.TRAIN
-    )
-    eval_dataloader = None
-    if pipeline_config.eval_input_path:
-        # pyre-ignore [16]
-        gl_cluster = train_dataloader.dataset.get_sampler_cluster()
-        eval_dataloader = create_dataloader(
-            data_config,
-            features,
-            pipeline_config.eval_input_path,
-            mode=Mode.EVAL,
-            gl_cluster=gl_cluster,
-        )
-
     ckpt_manager = checkpoint_util.CheckpointManager(
         pipeline_config.model_dir,
         keep_checkpoint_max=train_config.keep_checkpoint_max,
@@ -637,12 +621,33 @@ def train_and_evaluate(
                     "--continue_train)"
                 )
 
-    # Restore dataloader checkpoint state
+    # Restore dataloader checkpoint state before building the dataloader:
+    # create_dataloader eagerly starts persistent workers, which keep a
+    # fork-time copy of the dataset, so state set afterwards is invisible
+    # to them.
     dataloader_state: Optional[Dict[str, Any]] = None
     if ckpt_path and continue_train:
         dataloader_state = ckpt_manager.restore_dataloader_state(ckpt_path)
-        if dataloader_state:
-            train_dataloader.dataset.load_state_dict(dataloader_state)
+
+    # Build dataloader
+    train_dataloader = create_dataloader(
+        data_config,
+        features,
+        pipeline_config.train_input_path,
+        mode=Mode.TRAIN,
+        checkpoint_state=dataloader_state,
+    )
+    eval_dataloader = None
+    if pipeline_config.eval_input_path:
+        # pyre-ignore [16]
+        gl_cluster = train_dataloader.dataset.get_sampler_cluster()
+        eval_dataloader = create_dataloader(
+            data_config,
+            features,
+            pipeline_config.eval_input_path,
+            mode=Mode.EVAL,
+            gl_cluster=gl_cluster,
+        )
 
     sampler_type = _get_sampler_type(data_config)
 

tzrec/version.py

@@ -9,4 +9,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "1.2.18"
+__version__ = "1.2.19"