code: def test_concat_loader_with_ddp(
      use_distributed_sampler: bool,
      is_min_size_mode: bool,
      num_devices: Union[int, str],
  ):
      """Inspired by tests/trainer/test_supporters.py in lightning."""
      from nni.nas.oneshot.pytorch._dataloader import ConcatLoader
  
      mode = "min_size" if is_min_size_mode else "max_size_cycle"
      dim = 3
      n1 = 8
      n2 = 6
      n3 = 9
      dataloader = ConcatLoader(
          {
              "a": {
                  "a1": DataLoader(
                      RandomDataset(dim, n1), batch_size=1
                  ),
                  "a2": DataLoader(
                      RandomDataset(dim, n2), batch_size=1
                  ),
              },
              "b": DataLoader(RandomDataset(dim, n3), batch_size=1),
          }
      )
      print(len(dataloader))
      expected_length_before_ddp = n1 + n2 + n3
      assert len(dataloader) == expected_length_before_ddp
      model = BoringModel()
      trainer = Trainer(
          strategy="ddp",
          accelerator="cpu",
          devices=num_devices,
          use_distributed_sampler=use_distributed_sampler,
      )
      trainer.strategy.connect(model)
      trainer._data_connector.attach_data(
          model=model, train_dataloaders=dataloader
      )
      expected_length_after_ddp = (
          math.ceil(n3 / trainer.num_devices)
          + math.ceil(n1 / trainer.num_devices)
          + math.ceil(n2 / trainer.num_devices)
          if use_distributed_sampler
          else expected_length_before_ddp
      )
      print("Num devices =", trainer.num_devices)
      trainer.state.fn = "fit"
      trainer.state.stage = RunningStage.TRAINING
      trainer.fit_loop.setup_data()
      assert trainer.train_dataloader is not None
      assert trainer.num_training_batches == expected_length_after_ddp
  
imports: import math
  from typing import Union
  import pytest
  import torch
  import pytorch_lightning
  from pytorch_lightning import LightningModule, Trainer
  from pytorch_lightning.trainer.states import RunningStage
  from torch.utils.data import DataLoader, Dataset
label_desc: Use image node_id 3 for calling a global function with example usage: test_concat_loader_with_ddp(use_distributed_sampler, is_min_size_mode, num_devices) without return types