Repository: Lightning-AI/pytorch-lightning
File: tests/trainer/test_dataloaders.py (Python)

Code before:


@pytest.mark.skipif(platform.system() == 'Windows', reason='Does not apply to Windows platform.')
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')
def test_dataloader_distributed_sampler(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend """

    model = EvalModelTemplate()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=1,
        callbacks=[DistribSamplerCallback()]
    )
    trainer.fit(model)
    trainer.test(ckpt_path=None)


@pytest.mark.skipif(platform.system() == 'Windows', reason='Does not apply to Windows platform.')
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')
def test_dataloader_distributed_sampler_already_attached(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend when DistSampler already included on dataloader """

    class ModelWithDataLoaderDistributedSampler(EvalModelTemplate):

        def train_dataloader(self):
            dataloader = super().train_dataloader()
            dist_sampler = DistributedSampler(dataloader.dataset, shuffle=False)
            return DataLoader(
                dataset,
                batch_size=self.batch_size,
                drop_last=False,
                sampler=dist_sampler,
                shuffle=False
            )

    model = ModelWithDataLoaderDistributedSampler()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=100,
        callbacks=[DistribSamplerCallback()],
        replace_sampler_ddp=True,
    )
    trainer.fit(model)
    assert result == 1, "DDP Training failed"


Review comment (refactor):
hey @ananthsub 
Looks like you need to move this class outside the test function or we get the good old pickling error :) 

Code after:


@pytest.mark.skipif(platform.system() == 'Windows', reason='Does not apply to Windows platform.')
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')
def test_dataloader_distributed_sampler(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend """

    model = EvalModelTemplate()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=1,
        callbacks=[DistribSamplerCallback()]
    )
    trainer.fit(model)
    trainer.test(ckpt_path=None)


class ModelWithDataLoaderDistributedSampler(EvalModelTemplate):

    def train_dataloader(self):
        dataloader = super().train_dataloader()
        dist_sampler = DistributedSampler(dataloader.dataset, shuffle=True)
        return DataLoader(
            dataloader.dataset,
            batch_size=self.batch_size,
            drop_last=False,
            sampler=dist_sampler,
            shuffle=False
        )


@pytest.mark.skipif(platform.system() == 'Windows', reason='Does not apply to Windows platform.')
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='Test requires multiple GPUs')
def test_dataloader_distributed_sampler_already_attached(tmpdir):
    """ Test DistributedSampler and it's arguments for DDP backend when DistSampler already included on dataloader """

    model = ModelWithDataLoaderDistributedSampler()
    trainer = Trainer(
        gpus=[0, 1],
        num_nodes=1,
        distributed_backend='ddp_spawn',
        default_root_dir=tmpdir,
        max_steps=100,
        callbacks=[DistribSamplerCallback()],
        replace_sampler_ddp=True,
    )
    result = trainer.fit(model)
    assert result == 1, "DDP Training failed"