fix(fabric): forward set_epoch to underlying sampler in DistributedSamplerWrapper

[littlebullGit]

What does this PR do?

The DistributedSamplerWrapper now forwards set_epoch() calls to the original sampler if it supports the method.

Problem

When a custom sampler is wrapped by DistributedSamplerWrapper, calling set_epoch() on the wrapper does not propagate to the original sampler. This breaks samplers that rely on set_epoch() for epoch-dependent behavior like shuffling.

Solution

Override set_epoch() in DistributedSamplerWrapper to:

1. Call super().set_epoch(epoch) to set the epoch on the wrapper itself
2. Forward the call to the original sampler if it has a callable set_epoch method

This fix is generic and works for any sampler subclass that implements set_epoch(), not just specific implementations. It uses duck typing to check for the method's existence and callability.

Changes

• src/lightning/fabric/utilities/distributed.py: Added set_epoch() override to DistributedSamplerWrapper
• src/lightning/fabric/CHANGELOG.md: Added changelog entry
• tests/tests_fabric/utilities/test_distributed.py: Added test_distributed_sampler_wrapper_set_epoch with 100% branch coverage:
  • Case 1: Sampler WITH set_epoch method - verifies forwarding works
  • Case 2: Sampler WITHOUT set_epoch method - verifies graceful handling
  • Case 3: Sampler with non-callable set_epoch attribute - verifies it's skipped

Fixes #21454

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87%. Comparing base (aa0ee0d) to head (b27b18d).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #21456   +/-   ##
=======================================
  Coverage      87%      87%           
=======================================
  Files         270      270           
  Lines       24061    24067    +6     
=======================================
+ Hits        20857    20863    +6     
  Misses       3204     3204           

[deependujha]

Please update changelog too. :)

[littlebullGit]

Please update changelog too. :)

changelog is the main source of the merge conflict. So I am kind of hesitate to update it until someone reviewed. Not sure what is the best way to handle it. Updated.

[deependujha]

Yeah, but without it, we sometimes merge the PR without updates in changelog.

[littlebullGit]

Yeah, but without it, we sometimes merge the PR without updates in changelog.
Maybe we should add the changelog entry in the PR description and whoever merge it take it and add the the changelog ? just a thought.

[deependujha]

there're multiple maintainers maintaining this and other Lightning AI's OSS projects, I doubt if one will even remember, lol.

[bhimrazy]

Over all looks good.

[tchaton]

LGTM !