code: def test_distributed_sampler(setup_cluster):
      import torch
  
      data = mt.random.rand(1001, 32, dtype="f4")
      labels = mt.random.randint(0, 2, (1001, 10), dtype="f4")
  
      train_dataset = MarsDataset(data, labels)
  
      with pytest.raises(ValueError) as e:
          train_sampler = DistributedSampler(
              train_dataset, num_replicas=2, rank=-1
          )
      exec_msg = e.value.args[0]
      assert (
          exec_msg
          == "Invalid rank -1, rank should be in the interval [0, 1]"
      )
  
      train_sampler = DistributedSampler(
          train_dataset,
          num_replicas=2,
          rank=0,
          drop_last=True,
          shuffle=True,
      )
      assert len(train_sampler) == 500
      train_loader = torch.utils.data.DataLoader(
          dataset=train_dataset, batch_size=32, sampler=train_sampler
      )
      for _, (batch_data, batch_labels) in enumerate(train_loader):
          assert len(batch_data[0]) == 32
          assert len(batch_labels[0]) == 10
  
      train_sampler = DistributedSampler(
          train_dataset,
          num_replicas=2,
          rank=0,
          drop_last=False,
          shuffle=False,
      )
      train_sampler.set_epoch(10)
      assert len(train_sampler) == 501
      train_loader = torch.utils.data.DataLoader(
          dataset=train_dataset, batch_size=32, sampler=train_sampler
      )
      for _, (batch_data, batch_labels) in enumerate(train_loader):
          assert len(batch_data[0]) == 32
          assert len(batch_labels[0]) == 10
  
file_path: mars/mars/learn/contrib/pytorch/tests/test_dataset.py
imports: import pytest
  import os
  from .....None import tensor
  from .....None import dataframe
  from .....utils import lazy_import
  from ..None import MarsDataset, RandomSampler, SequentialSampler, SubsetRandomSampler, DistributedSampler, run_pytorch_script
label_desc: Use image node_id 5 for calling a global function with example usage: test_distributed_sampler(setup_cluster) without return types