Repository: Lightning-AI/pytorch-lightning
File: tests/accelerators/test_accelerator_connector.py (Python)

Code before:

    class CB(Callback):

        def on_fit_start(self, trainer, pl_module):
            assert trainer.use_ddp
            assert trainer.accelerator_connector.is_slurm_managing_tasks
            assert isinstance(trainer.accelerator, GPUAccelerator)
            assert isinstance(trainer.training_type_plugin, DDPPlugin)
            assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment)
            assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
            assert trainer.training_type_plugin.task_idx == 10
            raise SystemExit()

    model = BoringModel()
    trainer = Trainer(
        fast_dev_run=True,
        accelerator='ddp',
        gpus=2,
        callbacks=[CB()],
    )

    with pytest.raises(SystemExit):
        trainer.fit(model)


@SkipIf(min_gpus=1)
@mock.patch.dict(
    os.environ, {
        "CUDA_VISIBLE_DEVICES": "0,1",
        "SLURM_NTASKS": "2",
        "SLURM_JOB_NAME": "SOME_NAME",
        "SLURM_NODEID": "0",
        "LOCAL_RANK": "0",
        "SLURM_LOCALID": "10"
    }
)
@mock.patch('torch.cuda.device_count', return_value=2)
def test_accelerator_choice_ddp2_slurm(device_count_mock):

    class CB(Callback):

        def on_fit_start(self, trainer, pl_module):
            assert trainer.use_ddp2
            assert trainer.accelerator_connector.is_slurm_managing_tasks
            assert isinstance(trainer.accelerator, GPUAccelerator)
            assert isinstance(trainer.training_type_plugin, DDP2Plugin)
            assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment)
            assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
            assert trainer.training_type_plugin.task_idx == 10
            raise SystemExit()


Review comment (suggestion):
This name SkipIf, is it already decided?
Because the way you read it is not natural at all, and it expresses quite the opposite of what it should :) 
I think everyone who will see this the first time will scratch their head a few times 😄 

Code after:

    class CB(Callback):

        def on_fit_start(self, trainer, pl_module):
            assert trainer.use_ddp
            assert trainer.accelerator_connector.is_slurm_managing_tasks
            assert isinstance(trainer.accelerator, GPUAccelerator)
            assert isinstance(trainer.training_type_plugin, DDPPlugin)
            assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment)
            assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
            assert trainer.training_type_plugin.task_idx == 10
            raise SystemExit()

    model = BoringModel()
    trainer = Trainer(
        fast_dev_run=True,
        accelerator='ddp',
        gpus=2,
        callbacks=[CB()],
    )

    with pytest.raises(SystemExit):
        trainer.fit(model)


@RunIf(min_gpus=1)
@mock.patch.dict(
    os.environ, {
        "CUDA_VISIBLE_DEVICES": "0,1",
        "SLURM_NTASKS": "2",
        "SLURM_JOB_NAME": "SOME_NAME",
        "SLURM_NODEID": "0",
        "LOCAL_RANK": "0",
        "SLURM_LOCALID": "10"
    }
)
@mock.patch('torch.cuda.device_count', return_value=2)
def test_accelerator_choice_ddp2_slurm(device_count_mock):

    class CB(Callback):

        def on_fit_start(self, trainer, pl_module):
            assert trainer.use_ddp2
            assert trainer.accelerator_connector.is_slurm_managing_tasks
            assert isinstance(trainer.accelerator, GPUAccelerator)
            assert isinstance(trainer.training_type_plugin, DDP2Plugin)
            assert isinstance(trainer.training_type_plugin.cluster_environment, SLURMEnvironment)
            assert trainer.training_type_plugin.cluster_environment.local_rank() == 10
            assert trainer.training_type_plugin.task_idx == 10
            raise SystemExit()