Variable length number of batches is not supported

Nobody has claimed this yet.

### Bug description

BatchSamplers that do not always return same number of batches are broken in Pytorch Lightning.

Related: https://github.com/Lightning-AI/lightning/issues/17793.

There are several issues:
- `batch_size` is not required for a PyTorch `Sampler`--Pytorch Lightning is monkey patching this class such that a valid `Sampler` is incompatible with pytorch-lightning
- even though Pytorch Lightning repeatedly calls `len` on the Sampler, this information is disregarded, and the epoch often terminates early
- **if you use DDP** (>= 2 GPUs) and set the number of batches to a number larger than the actual number of batches (e.g. 100), **it causes a deadlock**
- behavior of DDP and single-gpu differs (see output of `batch_size`)

The code below reproduces all of these issues.

### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

```python
import os

import torch, numpy as np
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
from time import sleep

class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)

def __getitem__(self, index):
return self.data[index]

def __len__(self):
return self.len

class BoringModel(LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(32, 2)

def forward(self, x):
return self.layer(x)

def training_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("train_loss", loss)
return {"loss": loss}

def validation_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("valid_loss", loss)

def test_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("test_loss", loss)

def configure_optimizers(self):
return torch.optim.SGD(self.layer.parameters(), lr=0.1)

class VarLengthBatchSampler(torch.utils.data.sampler.Sampler):
def __init__(self, data_source, max_batch_size=None, drop_last=True, seed=12340,
batch_size=None):
"""Randomly return batches of length [1,max_batch_size].

Args:
data_source: a dataset
max_batch_size: the maximum batch size. This is required
drop_last: if True, drop the last batch if too small
seed: random seed
"""
super().__init__(data_source)
self.data_len = len(data_source)
self.max_batch_size = max_batch_size
self.seed = seed
# in actual application, need to increment this in the training loop
self.epoch = 0
# self.approx_len = int(np.ceil(self.data_len / (self.max_batch_size // 2)))
self.approx_len = 5 # to dramatically illustrate problem of not updating length
# self.approx_len = 100 # causes deadlock!!
print(f"approx_len: {self.approx_len}")
self.len = self.approx_len # pytorch lightning
self.drop_last = drop_last

def rand_bz(self):
return torch.randint(1,self.max_batch_size,(1,))[0]

def __iter__(self):
g = torch.Generator()
g.manual_seed(self.seed + self.epoch)
indices = torch.randperm(self.data_len, generator=g).tolist()
batch = [] # accumulate idxs, then reset
batches = [] # accumulate batches
bz = self.rand_bz()
# print(f"\n{bz=}")
for idx in indices:
batch.append(idx)
if len(batch) == bz:
batches.append(batch)
bz = self.rand_bz()
# print(f"\nnew {bz=}")
batch = []
if not self.drop_last and len(batch) > 0:
batches.append(batch)

# update length to match random sampling results
self.len = len(batches)
print(f"\nnumber of batches should be: {self.len}")
return iter(batches)

def __len__(self):
print(f"\n__len__: {self.len=}")
return self.len

class PLVarLengthBatchSampler(VarLengthBatchSampler):
def __init__(self, data_source, batch_size=None, max_batch_size=8, drop_last=True, seed=12340):
"""Wrapper for for pytorch lightning compatibility
"""
print(f"batch_size: {batch_size}")
super().__init__(data_source, max_batch_size=max_batch_size, drop_last=drop_last, seed=seed)
self.batch_size = max_batch_size

def run():
train_dset = RandomDataset(32, 64)
bz = 8
train_data = DataLoader(train_dset,
batch_sampler=PLVarLengthBatchSampler(train_dset, max_batch_size=bz))
val_data = DataLoader(RandomDataset(32, 64), batch_size=32)
test_data = DataLoader(RandomDataset(32, 64), batch_size=32)

model = BoringModel()
trainer = Trainer(
# devices=1, # uses first __len__ call for all lengths
devices=2, # uses approx_len for all lengths
default_root_dir=os.getcwd(),
num_sanity_val_steps=0,
max_epochs=5,
enable_model_summary=False
)
trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
trainer.test(model, dataloaders=test_data)

run()
```

### Error messages and logs

# One GPU
```
batch_size: None
approx_len: 5

number of batches should be: 15

__len__: self.len=15

__len__: self.len=15
Epoch 0: 0%| | 0[/15](https://file+.vscode-resource.vscode-cdn.net/15) [00:00
Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

### More info

_No response_

cc @justusschock @awaelchli

Start with the data connector entry point referenced in the report, site-packages/pytorch_lightning/trainer/connectors/data_connector.py, and reproduce the example with one GPU and with devices=2. Trace how Trainer.fit handles the custom batch sampler's changing __len__; done means variable batch counts complete without early termination or a DDP deadlock.

Written by the indexing model from the issue text.