docstr: 测试多卡下 set_dist_repro_dataloader 函数的执行结果是否正确
docstr_tok: ["\u6d4b\u8bd5\u591a\u5361\u4e0b","set_dist_repro_dataloader","\u51fd\u6570\u7684\u6267\u884c\u7ed3\u679c\u662f\u5426\u6b63\u786e"]
code: def check_set_dist_repro_dataloader(
      self, driver, dataloader, replaced_loader, shuffle
  ):
      """
      测试多卡下 set_dist_repro_dataloader 函数的执行结果是否正确
      """
      # 迭代两个 batch
      num_replicas = len(self.device)
      num_consumed_batches = 2
      already_seen_idx = set()
      if isinstance(
          replaced_loader.batch_sampler, BucketedBatchSampler
      ):
          sampler_states = replaced_loader.batch_sampler.set_epoch(4)
      else:
          sampler_states = (
              replaced_loader.batch_sampler.sampler.set_epoch(4)
          )
      for idx, batch in enumerate(replaced_loader):
          if idx >= num_consumed_batches:
              break
          already_seen_idx.update(batch.tolist())
      comm.barrier()
      if isinstance(
          replaced_loader.batch_sampler, BucketedBatchSampler
      ):
          sampler_states = replaced_loader.batch_sampler.state_dict()
      else:
          sampler_states = (
              replaced_loader.batch_sampler.sampler.state_dict()
          )
  
      # 重新加载，应该可以输出剩下的内容，且对于 OneflowNormalDataset 来说，排序后应该是一个 range
      left_idxes = set()
      if isinstance(
          replaced_loader.batch_sampler, BucketedBatchSampler
      ):
          batch_size = replaced_loader.batch_sampler.batch_size
          sampler_states["num_consumed_samples"] = (
              num_consumed_batches * batch_size * num_replicas
          )
          # 重新改造 dataloader
          new_loader = dataloader_with_bucketedbatchsampler(
              replaced_loader.dataset,
              length=replaced_loader.dataset._data,
              batch_size=batch_size,
              shuffle=shuffle,
              drop_last=False,
          )
          new_loader.batch_sampler.set_distributed(
              num_replicas=driver.world_size,
              rank=driver.global_rank,
              pad=True,
          )
          new_loader.batch_sampler.load_state_dict(sampler_states)
          new_loader.batch_sampler.set_epoch(4)
      else:
          batch_size = replaced_loader.batch_sampler.batch_size
          sampler_states["num_consumed_samples"] = (
              num_consumed_batches * batch_size * num_replicas
          )
          # 重新构造 dataloader
          new_loader = dataloader_with_randomsampler(
              replaced_loader.dataset,
              batch_size,
              shuffle,
              drop_last=False,
          )
          new_loader.batch_sampler.sampler.set_distributed(
              num_replicas=driver.world_size, rank=driver.global_rank
          )
          new_loader.batch_sampler.sampler.load_state_dict(
              sampler_states
          )
          new_loader.batch_sampler.sampler.set_epoch(4)
      for idx, batch in enumerate(new_loader):
          left_idxes.update(batch.tolist())
  
      assert (
          len(left_idxes) + len(already_seen_idx)
          == len(self.dataset) / num_replicas
      )
      assert (
          len(left_idxes | already_seen_idx)
          == len(self.dataset) / num_replicas
      )
  
imports: import os
  import sys
  import pytest
  from pathlib import Path
  from fastNLP.core.drivers.oneflow_driver.ddp import OneflowDDPDriver
  from fastNLP import prepare_oneflow_dataloader
  from fastNLP.core.samplers import RandomSampler, UnrepeatedSampler, BucketedBatchSampler, UnrepeatedRandomSampler, UnrepeatedSequentialSampler
  from tests.helpers.models.oneflow_model import OneflowNormalModel_Classification_1
  from tests.helpers.datasets.oneflow_data import OneflowNormalDataset, OneflowNormalXYDataset
  from tests.helpers.utils import recover_logger
  from fastNLP.envs.distributed import rank_zero_rm
  from fastNLP import logger
  from fastNLP.core.drivers.oneflow_driver.dist_utils import fastnlp_oneflow_all_gather
  from fastNLP.envs.imports import _NEED_IMPORT_ONEFLOW
label_desc: Use image node_id 16 for calling the TestSetDistReproDataloader obj's underlying member method code with example usage: obj.check_set_dist_repro_dataloader(driver, dataloader, replaced_loader, shuffle) without return types