ronantakizawa/codereview-bench huggingface dataset:
Repository: Lightning-AI/pytorch-lightning
File: tests/base/boring_model.py (Python)

Code before:

    def test_dataloader(self):
        return DataLoader(RandomDataset(32, 64))


class BoringDataModule(LightningDataModule):
    def __init__(self, data_dir: str = './'):
        super().__init__()
        self.data_dir = data_dir
        self.non_picklable = None
        self.checkpoint_state: Optional[str] = None

    def prepare_data(self):
        self.random_full = RandomDataset(32, 192)

    def setup(self, stage: Optional[str] = None):
        if stage == "fit" or stage is None:
            self.random_train = Subset(self.random_full, indices=range(64))
            self.random_val = Subset(self.random_full, indices=range(64, 128))
            self.dims = self.random_train[0].shape

        if stage == "test" or stage is None:
            self.random_test = Subset(self.random_full, indices=range(128, 192))
            self.dims = getattr(self, "dims", self.random_test[0].shape)

        self.non_picklable = lambda x: x ** 2

    def train_dataloader(self):
        return DataLoader(self.random_train)

    def val_dataloader(self):
        return DataLoader(self.random_val)

    def test_dataloader(self):
        return DataLoader(self.random_test)

    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        checkpoint[self.__class__.__name__] = self.__class__.__name__

    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        self.checkpoint_state = checkpoint.get(self.__class__.__name__)


Review comment (suggestion):
I think the tests that require `non_picklable` should subclass and override setup(), just as we do for BoringModel

Code after:

    def test_dataloader(self):
        return DataLoader(RandomDataset(32, 64))


class BoringDataModule(LightningDataModule):
    def __init__(self, data_dir: str = './'):
        super().__init__()
        self.data_dir = data_dir
        self.non_picklable = None
        self.checkpoint_state: Optional[str] = None

    def prepare_data(self):
        self.random_full = RandomDataset(32, 192)

    def setup(self, stage: Optional[str] = None):
        if stage == "fit" or stage is None:
            self.random_train = Subset(self.random_full, indices=range(64))
            self.random_val = Subset(self.random_full, indices=range(64, 128))
            self.dims = self.random_train[0].shape

        if stage == "test" or stage is None:
            self.random_test = Subset(self.random_full, indices=range(128, 192))
            self.dims = getattr(self, "dims", self.random_test[0].shape)

    def train_dataloader(self):
        return DataLoader(self.random_train)

    def val_dataloader(self):
        return DataLoader(self.random_val)

    def test_dataloader(self):
        return DataLoader(self.random_test)
Repository: Lightning-AI/pytorch-lightning
File: tests/core/test_datamodules.py (Python)

Code before:
        weights_summary=None,
    )
    trainer.test(model, datamodule=dm)


def test_full_loop(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert result == 1

    # test
    result = trainer.test(datamodule=dm)
    # result = result[0]
    # assert result['test_acc'] > 0.8


def test_trainer_attached_to_dm(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        limit_train_batches=2,
        limit_val_batches=2,
        limit_test_batches=2,
        weights_summary=None,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert result == 1
    assert dm.trainer is not None

    # test
    result = trainer.test(datamodule=dm)

Review comment (suggestion):
```suggestion
    # TODO: add end-to-end test for: assert result[0]['test_acc'] > 0.8
```

Code after:
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
    )
    trainer.test(model, datamodule=dm)


def test_full_loop(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
    assert result

    # test
    result = trainer.test(datamodule=dm)
    # TODO: add end-to-end test
    # assert result[0]['test_acc'] > 0.8


def test_trainer_attached_to_dm(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        limit_train_batches=2,
        limit_val_batches=2,
        limit_test_batches=2,
        weights_summary=None,
        deterministic=True,
    )

    # fit model
Repository: Lightning-AI/pytorch-lightning
File: tests/core/test_datamodules.py (Python)

Code before:
    assert dm.trainer is not None


@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
def test_full_loop_single_gpu(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        gpus=1,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert result

    # test
    result = trainer.test(datamodule=dm)
    # result = result[0]
    # assert result['test_acc'] > 0.8


@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
def test_full_loop_dp(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        accelerator='dp',
        gpus=2,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert result

    # test
    result = trainer.test(datamodule=dm)
    # result = result[0]

Review comment (suggestion):
```suggestion
    # TODO: Add end-to-end test for: assert result[0]['test_acc'] > 0.8
```

Code after:
    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
    assert dm.trainer is not None

    # test
    result = trainer.test(datamodule=dm)
    result = result[0]
    assert dm.trainer is not None


@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
def test_full_loop_single_gpu(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        gpus=1,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
    assert result

    # test
    result = trainer.test(datamodule=dm)
    # TODO: add end-to-end test
    # assert result[0]['test_acc'] > 0.8


@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
def test_full_loop_dp(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        accelerator='dp',
        gpus=2,
        deterministic=True,
    )
Repository: Lightning-AI/pytorch-lightning
File: tests/core/test_datamodules.py (Python)

Code before:


@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
def test_full_loop_dp(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        accelerator='dp',
        gpus=2,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert result

    # test
    result = trainer.test(datamodule=dm)
    # result = result[0]
    # assert result['test_acc'] > 0.8


@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
def test_dm_transfer_batch_to_device(tmpdir):
    class CustomBatch:
        def __init__(self, data):
            self.samples = data[0]
            self.targets = data[1]

    class CurrentTestDM(LightningDataModule):

        hook_called = False

        def transfer_batch_to_device(self, data, device):
            self.hook_called = True
            data.samples = data.samples.to(device)
            data.targets = data.targets.to(device)
            return data

    dm = CurrentTestDM()
    model = BoringModel()

    batch = CustomBatch((torch.zeros(5, 32), torch.ones(5, 1, dtype=torch.long)))

    trainer = Trainer(gpus=1)

Review comment (suggestion):
```suggestion
    # TODO: Add end-to-end test for: assert result[0]['test_acc'] > 0.8
```

Code after:
    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
    assert result

    # test
    result = trainer.test(datamodule=dm)
    # TODO: add end-to-end test
    # assert result[0]['test_acc'] > 0.8


@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
def test_full_loop_dp(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
        accelerator='dp',
        gpus=2,
        deterministic=True,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
    assert result

    # test
    result = trainer.test(datamodule=dm)
    # TODO: add end-to-end test
    # assert result[0]['test_acc'] > 0.8


@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
def test_dm_transfer_batch_to_device(tmpdir):
    class CustomBatch:
        def __init__(self, data):
            self.samples = data[0]
            self.targets = data[1]

    class CurrentTestDM(LightningDataModule):

        hook_called = False

        def transfer_batch_to_device(self, data, device):
            self.hook_called = True
            data.samples = data.samples.to(device)
            data.targets = data.targets.to(device)
Repository: Lightning-AI/pytorch-lightning
File: tests/core/test_datamodules.py (Python)

Code before:
def test_train_loop_only(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    model.validation_step = None
    model.validation_step_end = None
    model.validation_epoch_end = None
    model.test_step = None
    model.test_step_end = None
    model.test_epoch_end = None

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert result
    # assert trainer.callback_metrics['loss'] < 0.6


def test_train_val_loop_only(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    model.validation_step = None
    model.validation_step_end = None
    model.validation_epoch_end = None

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert result
    # assert trainer.callback_metrics['train_loss'] < 0.6


def test_dm_checkpoint_save(tmpdir):
    class CustomBoringModel(BoringModel):
        def validation_step(self, batch, batch_idx):
            out = super().validation_step(batch, batch_idx)

Review comment (suggestion):
```suggestion
    # TODO: Add end-to-end test for:  assert trainer.callback_metrics['loss'] < 0.6
```

Code after:

def test_train_loop_only(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    model.validation_step = None
    model.validation_step_end = None
    model.validation_epoch_end = None
    model.test_step = None
    model.test_step_end = None
    model.test_epoch_end = None

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
    assert result
    # TODO: add end-to-end test
    # assert trainer.callback_metrics['loss'] < 0.6


def test_train_val_loop_only(tmpdir):
    reset_seed()

    dm = BoringDataModule()
    model = BoringModel()

    model.validation_step = None
    model.validation_step_end = None
    model.validation_epoch_end = None

    trainer = Trainer(
        default_root_dir=tmpdir,
        max_epochs=1,
        weights_summary=None,
    )

    # fit model
    result = trainer.fit(model, dm)
    assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
    assert result
    # TODO: add end-to-end test
    # assert trainer.callback_metrics['train_loss'] < 0.6
Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/supporters.py (Python)

Code before:

class CombinedLoaderIterator(object):
    """
    Custom Iterator returning data from multple loaders, and allows sampling in parallel
    """
    def __init__(self, loaders: Any):
        """

        Args:
            loaders: the loaders to sample from. Can be all kind of collection

        """
        self.loaders = loaders
        self._loader_iters = None

    @property
    def loader_iters(self) -> Any:
        """
        Get the `_loader_iters` and create one if it is None.
        """
        if self._loader_iters is None:
            self._loader_iters = self.create_loader_iters(self.loaders)

        return self._loader_iters

    def __iter__(self) -> Any:
        return self

    def __next__(self) -> Any:
        """
        Fetches the next batch from multiple data loaders

        Returns:
            a collections of batch data

        """
        return self.request_next_batch(self.loader_iters)

    @staticmethod
    def request_next_batch(loader_iters: Union[Iterator, Sequence, Mapping]) -> Any:
        """
        Return the batch of data from multiple iterators.

        Args:
            loader_iters: a collections of iterators

        Returns
            Any: a collections of batch data

        """
        return apply_to_collection(loader_iters, Iterator, next)

Review comment (performance):
```suggestion
            a collections of batch data
```

Code after:
            loaders: the loaders to sample from. Can be all kind of collection

        """
        self.loaders = loaders
        self._loader_iters = None

    @property
    def loader_iters(self) -> Any:
        """
        Get the `_loader_iters` and create one if it is None.
        """
        if self._loader_iters is None:
            self._loader_iters = self.create_loader_iters(self.loaders)

        return self._loader_iters

    def __iter__(self) -> Any:
        return self

    def __next__(self) -> Any:
        """
        Fetches the next batch from multiple data loaders

        Returns:
            a collections of batch data

        """
        return self.request_next_batch(self.loader_iters)

    @staticmethod
    def request_next_batch(loader_iters: Union[Iterator, Sequence, Mapping]) -> Any:
        """
        Return the batch of data from multiple iterators.

        Args:
            loader_iters: a collections of iterators

        Returns
            Any: a collections of batch data

        """
        return apply_to_collection(loader_iters, Iterator, next)

    @staticmethod
    def create_loader_iters(
        loaders: Union[Any, Iterator, Sequence, Mapping]
    ) -> Union[Any, Iterator, Sequence, Mapping]:
        """
        Create and return a collection of iterators from loaders.

        Args:
Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/utilities/imports.py (Python)

Code before:
        # Python 3.7+
        return False


def _get_version(package: str) -> LooseVersion:
    return LooseVersion(pkg_resources.get_distribution(package).version)


_IS_WINDOWS = platform.system() == "Windows"

_APEX_AVAILABLE = _module_available("apex.amp")
_BOLTS_AVAILABLE = _module_available('pl_bolts')
_FAIRSCALE_AVAILABLE = not _IS_WINDOWS and _module_available('fairscale.nn.data_parallel')
_TORCH_GREATER_EQUAL_1_6 = _get_version("torch") >= LooseVersion("1.6.0")
_FAIRSCALE_PIPE_AVAILABLE = (
    _FAIRSCALE_AVAILABLE and _TORCH_GREATER_EQUAL_1_6 and _get_version('fairscale') <= LooseVersion("0.1.3")
)
_GROUP_AVAILABLE = not _IS_WINDOWS and _module_available('torch.distributed.group')
_HOROVOD_AVAILABLE = _module_available("horovod.torch")
_HYDRA_AVAILABLE = _module_available("hydra")
_HYDRA_EXPERIMENTAL_AVAILABLE = _module_available("hydra.experimental")
_NATIVE_AMP_AVAILABLE = _module_available("torch.cuda.amp") and hasattr(torch.cuda.amp, "autocast")
_OMEGACONF_AVAILABLE = _module_available("omegaconf")
_PYTORCH_PRUNE_AVAILABLE = _module_available('torch.nn.utils.prune')
_RPC_AVAILABLE = not _IS_WINDOWS and _module_available('torch.distributed.rpc')
_TORCHTEXT_AVAILABLE = _module_available("torchtext")
_TORCHVISION_AVAILABLE = _module_available('torchvision')
_XLA_AVAILABLE = _module_available("torch_xla")


Review comment (suggestion):
```suggestion
_RPC_AVAILABLE = not _IS_WINDOWS and _module_available('torch.distributed.rpc')
_TORCH_GREATER_EQUAL_1_6 = _get_version("torch") >= LooseVersion("1.6.0")
_TORCHTEXT_AVAILABLE = _module_available("torchtext")
```

Code after:
        # Python 3.7+
        return False


def _get_version(package: str) -> LooseVersion:
    return LooseVersion(pkg_resources.get_distribution(package).version)


_IS_WINDOWS = platform.system() == "Windows"
_TORCH_GREATER_EQUAL_1_6 = _get_version("torch") >= LooseVersion("1.6.0")

_APEX_AVAILABLE = _module_available("apex.amp")
_BOLTS_AVAILABLE = _module_available('pl_bolts')
_FAIRSCALE_AVAILABLE = not _IS_WINDOWS and _module_available('fairscale.nn.data_parallel')
_FAIRSCALE_PIPE_AVAILABLE = (
    _FAIRSCALE_AVAILABLE and _TORCH_GREATER_EQUAL_1_6 and _get_version('fairscale') <= LooseVersion("0.1.3")
)
_GROUP_AVAILABLE = not _IS_WINDOWS and _module_available('torch.distributed.group')
_HOROVOD_AVAILABLE = _module_available("horovod.torch")
_HYDRA_AVAILABLE = _module_available("hydra")
_HYDRA_EXPERIMENTAL_AVAILABLE = _module_available("hydra.experimental")
_NATIVE_AMP_AVAILABLE = _module_available("torch.cuda.amp") and hasattr(torch.cuda.amp, "autocast")
_OMEGACONF_AVAILABLE = _module_available("omegaconf")
_RPC_AVAILABLE = not _IS_WINDOWS and _module_available('torch.distributed.rpc')
_TORCHTEXT_AVAILABLE = _module_available("torchtext")
_TORCHVISION_AVAILABLE = _module_available('torchvision')
_XLA_AVAILABLE = _module_available("torch_xla")
Repository: Lightning-AI/pytorch-lightning
File: tests/helpers/skipif.py (Python)

Code before:
            ...

    >>> from pprint import pprint
    >>> pprint(create_skipif(min_torch="99"))
    {'condition': True, 'reason': 'test requires minimal version `torch>=99'}
    >>> pprint(create_skipif(min_torch="0.0"))
    {'condition': False, 'reason': 'no reason, just go test it...'}
    """
    conditions = []
    reasons = []

    if min_gpus:
        conditions.append(torch.cuda.device_count() < 2)
        reasons.append(f"multi-GPU machine with at least {min_gpus}")

    if min_torch:
        torch_version = LooseVersion(get_distribution("torch").version)
        conditions.append(torch_version < LooseVersion(min_torch))
        reasons.append(f"minimal version `torch>={min_torch}")

    if pt_quant:
        conditions.append(not _TORCH_QUANTIZE_AVAILABLE or _MISS_QUANT_DEFAULT)
        reasons.append("PyTorch quantization")

    if not any(conditions):
        return dict(condition=False, reason="no reason, just go test it...")

    reasons = [rs for cond, rs in zip(conditions, reasons) if cond]
    reason = "test requires " + ' + '.join(reasons)
    return dict(condition=any(conditions), reason=reason)


@pytest.mark.skipif(**create_skipif(min_torch="99"))
def test_always_skip():
    exit(1)


@pytest.mark.skipif(**create_skipif(min_torch="0.0"))
def test_always_pass():
    assert True


Review comment (suggestion):
```suggestion
        return dict(condition=False, reason="Conditions satisfied, going ahead with the test")
```

Code after:
    >>> from pprint import pprint
    >>> pprint(skipif_args(min_torch="99", min_gpus=0))
    {'condition': True, 'reason': 'Required: [torch>=99]'}
    >>> pprint(skipif_args(min_torch="0.0", min_gpus=0))  # doctest: +NORMALIZE_WHITESPACE
    {'condition': False, 'reason': 'Conditions satisfied, going ahead with the test.'}
    """
    conditions = []
    reasons = []

    if min_gpus:
        conditions.append(torch.cuda.device_count() < min_gpus)
        reasons.append(f"GPUs>={min_gpus}")

    if min_torch:
        torch_version = LooseVersion(get_distribution("torch").version)
        conditions.append(torch_version < LooseVersion(min_torch))
        reasons.append(f"torch>={min_torch}")

    if quant_available:
        _miss_default = 'fbgemm' not in torch.backends.quantized.supported_engines
        conditions.append(not _TORCH_QUANTIZE_AVAILABLE or _miss_default)
        reasons.append("PyTorch quantization is available")

    if not any(conditions):
        return dict(condition=False, reason="Conditions satisfied, going ahead with the test.")

    reasons = [rs for cond, rs in zip(conditions, reasons) if cond]
    return dict(condition=any(conditions), reason=f"Required: [{' + '.join(reasons)}]",)


@pytest.mark.skipif(**skipif_args(min_torch="99"))
def test_always_skip():
    exit(1)


@pytest.mark.skipif(**skipif_args(min_torch="0.0"))
def test_always_pass():
    assert True