code: def test_exception_early_stop_asap(setup_rpc):
      """Even the first partitions have finished to process, the partition before
      the failed partition should be killed as soon as possible.
      """
  
      class ExpectedException(Exception):
          pass
  
      class Pass(nn.Module):
          def forward(self, x):
              return x
  
      counter = 0
  
      class Counter(nn.Module):
          def forward(self, x):
              time.sleep(0.1)
  
              nonlocal counter
              counter += 1
  
              return x
  
      class Raise(nn.Module):
          def forward(self, x):
              raise ExpectedException()
  
      model = nn.Sequential(Pass(), Pass(), Counter(), Raise())
      model = Pipe(model, chunks=3)
  
      with pytest.raises(ExpectedException):
          model(torch.rand(3))
  
      # If the early stop doesn't work, it would be 3 instead.
      assert counter == 2
  
imports: from collections import OrderedDict
  from copy import deepcopy
  import time
  import pytest
  import random
  import torch
  from torch import nn
  from torch import Tensor
  from torch.distributed.pipeline.sync import Pipe, NoChunk, WithDevice
  from torch.distributed.pipeline.sync.pipe import PipeSequential
  from torch.testing._internal.common_utils import run_tests, TEST_CUDA
  from torch.testing._internal.common_cuda import TEST_MULTIGPU
label_desc: Use image node_id 15 for calling a global function with example usage: test_exception_early_stop_asap(setup_rpc) and returns: x, x