Refactor bucket batch sampler to support DDP sampling

[guoqing-zhou]

(NOTE: delete these instructional lines as you fill-out this PR template)

(NOTE: _template is designed to be filled-in and used as the squashed commit message for the entire PR. Italicized text is intended to be deleted as you fill in this template. Use the text between the ---)

---

High level summary of changes. Try to keep this as short and informative as possible: less is more.

Describe your changes. You can be more detailed and descriptive here. If it is a code change, Be sure to answer:

• What is changing?
• What is the new or fixed functionality?
• Why or when would someone want to use these changes?
• How can someone use these changes?

---

Summary

High level summary of changes. Try to keep this as short and informative as possible: less is more.

Details

Describe your changes. You can be more detailed and descriptive here.

Usage

How does a user interact with the changed code?

python -m your.new.module -and -all -options

Testing

How do you prove that your code behaves the way you claim?

Tests for these changes can be run via:

pytest -v tests/your/new/or/existing/test_functions.py::test_function

(NOTE: also delete this checklist as you fill-out this PR template)

Most of the changes to files with extensions *.py, *.yaml, *.yml, Dockerfile* or requirements.txt DO REQUIRE both pytest- and jet- CI stages.

• Did you review the Before your PR is "Ready for review" section before asking for review?
• Did you make sure your changes have tests? Did you test your changes locally?
• Can you add the SKIP_CI label to your PR?
• Can you add the PYTEST_NOT_REQUIRED label to your PR?
• Can you add the JET_NOT_REQUIRED label to your PR?

[DejunL]

Did my first pass. Have some minor comments. I would need to look more into the how lightning engage with the DistributedSampler to get an understanding of the underlying logics so to come up with a test