code: def test_torch_fail_on_nccl_timeout(ray_start_4_cpus_2_gpus):
      """Tests that TorchTrainer raises exception on NCCL timeouts."""
  
      def train_fn():
          model = torch.nn.Linear(1, 1)
          model = train.torch.prepare_model(model)
  
          # Rank 0 worker will never reach the collective operation.
          # NCCL should timeout.
          if train.get_context().get_world_rank() == 0:
              while True:
                  time.sleep(100)
  
          torch.distributed.barrier()
  
      trainer = TorchTrainer(
          train_fn,
          scaling_config=ScalingConfig(num_workers=2, use_gpu=True),
          torch_config=TorchConfig(timeout_s=5),
      )
  
      # Training should fail and not hang.
      with pytest.raises(TrainingFailedError) as exc_info:
          trainer.fit()
      assert isinstance(exc_info.value.__cause__, RayTaskError)
  
imports: import json
  import os
  import time
  from pathlib import Path
  from typing import Dict, List, Union
  from unittest.mock import patch
  import pytest
  import torch
  import torchvision
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import DataLoader, DistributedSampler
  import ray
  import ray.data
  from ray import train
  from ray.exceptions import RayTaskError
  from ray.train import ScalingConfig
  from ray.train._internal.worker_group import WorkerGroup
  from ray.train.constants import DEFAULT_NCCL_SOCKET_IFNAME
  from ray.train.examples.pytorch.torch_linear_example import LinearDataset
  from ray.train.torch.config import TorchConfig, _TorchBackend
  from ray.train.torch.torch_trainer import TorchTrainer
  from ray.train.trainer import TrainingFailedError
label_desc: Use image node_id 10 for calling a global function with example usage: test_torch_fail_on_nccl_timeout(ray_start_4_cpus_2_gpus) without return types