def METHOD_NAME(self):
    splitter = TextDatasetSplitter(
        dataset_name="test",
        dataset_size=10000,
        shard_size=100,
        num_epochs=1,
    )
    ds_manager = BatchDatasetManager(TaskType.TRAINING, 10, splitter)
    worker_id = 0
    task = ds_manager.get_task(NodeType.WORKER, worker_id)
    self.assertEqual(task.task_id, 0)
    self.assertEqual(len(ds_manager.todo), 99)
    self.assertEqual(len(ds_manager.doing), 1)
    self.assertFalse(ds_manager.completed())
    checkpoint = ds_manager.checkpoint()
    self.assertListEqual(checkpoint.todo[0], [100, 200])
    ds_manager.restore_checkpoint(checkpoint)
    self.assertEqual(ds_manager.todo[0].shard.start, 0)
    self.assertEqual(ds_manager.todo[0].shard.end, 100)
    self.assertEqual(ds_manager.todo[0].shard.record_indices, None)