class_docstr: Wrap the PyTorch DataLoader WorkerLoop to perform caching and profiling.
class_docstr_tok: ["Wrap","the","PyTorch","DataLoader","WorkerLoop","to","perform","caching","and","profiling","."]
code: def __call__(
      self,
      dataset_kind: Any,
      dataset: Any,
      index_queue: Any,
      data_queue: Any,
      done_event: Any,
      auto_collation: Any,
      collate_fn: Any,
      drop_last: Any,
      base_seed: Any,
      init_fn: Any,
      worker_id: Any,
      *args: Any,
      **kwargs: Any,
  ) -> None:
      from torch.utils.data._utils import worker
  
      from lightning.data.streaming.cache import Cache
  
      enable_profiling = (
          self._global_rank == 0
          and worker_id == 0
          and _VIZ_TRACKER_AVAILABLE
          and self._profile
      )
  
      if enable_profiling:
          from viztracer import VizTracer
  
          tracer = VizTracer(
              output_file=os.path.join(os.getcwd(), "trace.json")
          )
          tracer.start()
  
      # Reload to remove the patching
      reloaded_worker = reload(worker)
      create_fetcher = _DatasetKind.create_fetcher
      fetcher = None
  
      def create_fetcher_fn(
          *args: Any, **kwargs: Any
      ) -> "_BaseDatasetFetcher":
          nonlocal fetcher
          fetcher = create_fetcher(*args, **kwargs)
          return fetcher
  
      _DatasetKind.create_fetcher = create_fetcher_fn  # type: ignore
  
      reloaded_worker._worker_loop(
          dataset_kind,
          dataset,
          index_queue,
          data_queue,
          done_event,
          auto_collation,
          collate_fn,
          drop_last,
          base_seed,
          init_fn,
          worker_id,
          *args,
          **kwargs,
      )
  
      if dataset_kind == _DatasetKind.Map:
          assert fetcher
          for v in fetcher.dataset.__dict__.values():
              if isinstance(v, Cache):
                  v.done()
  
      if enable_profiling:
          tracer.stop()
          tracer.save()
  
imports: import asyncio
  import inspect
  import logging
  import os
  from importlib import reload
  from typing import Any, Callable, Dict, List, Optional, Union
  import torch
  from torch.utils.data import Dataset, IterableDataset
  from torch.utils.data._utils.collate import default_collate
  from torch.utils.data._utils.fetch import _BaseDatasetFetcher
  from torch.utils.data.dataloader import DataLoader, _BaseDataLoaderIter, _DatasetKind, _MultiProcessingDataLoaderIter, _SingleProcessDataLoaderIter
  from torch.utils.data.sampler import BatchSampler, Sampler
  from lightning.data.streaming import Cache
  from lightning.data.streaming.combined import CombinedStreamingDataset
  from lightning.data.streaming.constants import _DEFAULT_CHUNK_BYTES, _TORCH_GREATER_EQUAL_2_1_0, _VIZ_TRACKER_AVAILABLE
  from lightning.data.streaming.dataset import StreamingDataset
  from lightning.data.streaming.sampler import CacheBatchSampler
  from lightning.data.utilities.env import _DistributedEnv
label_desc: Use image node_id 2 for calling the WorkerLoop obj's underlying member method code with example usage: obj.__call__(dataset_kind, dataset, index_queue, data_queue, done_event, auto_collation, collate_fn, drop_last, base_seed, init_fn, worker_id) and returns: fetcher