feat:Add dataset truncation to prevent deadlocks during training

[ved1beta]

What does this PR do?

• Add truncate_dataset and gradient_steps parameters to BatchSamplerShard
• Implement optimal dataset size calculation
• Add truncation logic to both split and non-split batch iterators Fixes #2019

Before submitting

• [X] Was this discussed/approved via a Github issue or the Please add a link to it if that's the case.
• [X] Did you write any new necessary tests?

Who can review?

@SunMarc @muellerzr