docstr: The CacheBatchSampler handles the generation of batch indices.
  
  If the cache isn't filled, the batch sampler alternates with ordered indices for the writer to chunk the dataset
  If the cache is filled, it acts as normal BatchSampler.
  
  Arguments:
      dataset_size: The size of the dataset.
      num_replicas: The number of processes involves in the distributed training.
      global_rank: The global_rank of the given process
      num_workers: The number of workers provided to the DataLoader.
      batch_size: The number of items in a batch.
      drop_last: Whether to drop the last batch of data.
      shuffle: Whether the data should be shuffled.
      cache: The cache associated to the dataset.
docstr_tok: ["The","CacheBatchSampler","handles","the","generation","of","batch","indices",".","If","the","cache","is","n't","filled",",","the","batch","sampler","alternates","with","ordered","indices","for","the","writer","to","chunk","the","dataset","If","the","cache","is","filled",",","it","acts","as","normal","BatchSampler",".","Arguments",":","dataset_size",":","The","size","of","the","dataset",".","num_replicas",":","The","number","of","processes","involves","in","the","distributed","training",".","global_rank",":","The","global_rank","of","the","given","process","num_workers",":","The","number","of","workers","provided","to","the","DataLoader",".","batch_size",":","The","number","of","items","in","a","batch",".","drop_last",":","Whether","to","drop","the","last","batch","of","data",".","shuffle",":","Whether","the","data","should","be","shuffled",".","cache",":","The","cache","associated","to","the","dataset","."]
code: def __init__(
      self,
      dataset_size: int,
      num_replicas: int,
      global_rank: int,
      num_workers: int,
      batch_size: int,
      drop_last: bool,
      shuffle: bool,
      cache: Any,
  ):
      """The CacheBatchSampler handles the generation of batch indices.
  
      If the cache isn't filled, the batch sampler alternates with ordered indices for the writer to chunk the dataset
      If the cache is filled, it acts as normal BatchSampler.
  
      Arguments:
          dataset_size: The size of the dataset.
          num_replicas: The number of processes involves in the distributed training.
          global_rank: The global_rank of the given process
          num_workers: The number of workers provided to the DataLoader.
          batch_size: The number of items in a batch.
          drop_last: Whether to drop the last batch of data.
          shuffle: Whether the data should be shuffled.
          cache: The cache associated to the dataset.
  
      """
      self._dataset_size = dataset_size
      self._num_replicas = num_replicas
      self._global_rank = global_rank
      self._cache = cache
      self._shuffle = shuffle
      self._num_workers = num_workers or 1
      self._shuffled_chunk_intervals = None
      self._batch_size = batch_size
  
      self._drop_last = drop_last
      self._length = 0
  
      # Before starting, ensures the chunk indices are properly defined.
      self._validate()
  
file_path: pytorch-lightning/src/lightning/data/streaming/sampler.py
imports: import logging
  from dataclasses import dataclass
  from typing import Any, Dict, Iterator, List, Optional, Union
  import numpy
label_desc: Use image node_id 1 to create a new CacheBatchSampler object  with example: obj = CacheBatchSampler(dataset_size, num_replicas, global_rank, num_workers, batch_size, drop_last, shuffle, cache)