Lightning-AI / Lightning-AI/torchmetrics
`Metric.sync_context` leaves the metric synchronized when `compute()` raises
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.5k
- Forks
- 526
- Avg merge
- 6d 11h
- Merged PRs (30d)
- 5
Description
## Bug description
After distributed synchronization succeeds, an exception raised by a metric's
`compute()` skips the local-state restoration at the end of `sync_context`.
The metric remains synchronized. A subsequent `update()` followed by `compute()`
fails with `TorchMetricsUserError: The Metric has already been synced.`
The initial exception in this example is intentional: a retrieval query has no
positive target and `empty_target_action="error"`. After adding a positive target,
the accumulated query is valid, but the leftover synchronization state prevents
its evaluation.
## Minimal reproduction
Save the following as `repro_compute.py` and run on CPU:
```bash
python -m torch.distributed.run --standalone --nproc_per_node=2 repro_compute.py
```
```python
"""Run: python -m torch.distributed.run --standalone --nproc_per_node=2 repro_compute.py"""
from datetime import timedelta
import json
import torch
import torch.distributed as dist
from torchmetrics.retrieval import RetrievalMRR
from torchmetrics.utilities.exceptions import TorchMetricsUserError
def main():
torch.set_num_threads(1)
dist.init_process_group("gloo", timeout=timedelta(seconds=30))
try:
rank = dist.get_rank()
metric = RetrievalMRR(empty_target_action="error", compute_with_cache=False)
metric.update(
torch.tensor([0.95 - 0.1 * rank]), torch.tensor([0]), indexes=torch.tensor([0])
)
first_error = None
try:
metric.compute()
except ValueError as error:
first_error = str(error)
synced_after_error = metric._is_synced
metric.update(
torch.tensor([0.55 - 0.1 * rank]), torch.tensor([1]), indexes=torch.tensor([0])
)
try:
result = float(metric.compute())
second_error = None
except TorchMetricsUserError as error:
result = None
second_error = str(error)
record = {
"rank": rank,
"first_error": first_error,
"synced_after_error": synced_after_error,
"mrr": result,
"second_error": second_error,
}
records = [None] * dist.get_world_size()
dist.all_gather_object(records, record)
if rank == 0:
print("RESULT_JSON=" + json.dumps(records), flush=True)
finally:
dist.destroy_process_group()
if __name__ == "__main__":
main()
```
## Actual and expected behavior
On both ranks, the initial error is `no positive target`, `synced_after_error`
is `true`, and the second call fails with `The Metric has already been synced.`
The initial error should still propagate. When `should_unsync=True`, exiting
the synchronized context should restore the cached local states, including on
this exceptional path. After the common continuation, the global MRR should be
approximately `0.33333334` (two higher-scored negatives precede the positives).
Calling `reset()` is not equivalent: it discards accepted observations.
## Suggested fix
Keep `self.sync(...)` in its current location, and move the existing `unsync`
call into a `finally` block around `yield`:
```python
try:
yield
finally:
self.unsync(should_unsync=self._is_synced and should_unsync)
```
This preserves `should_unsync=False` and does not add a collective. It covers
exceptions in the yielded body after synchronization succeeds, not failed
collectives, process termination, or arbitrary user mutation of metric state.
I have prepared a small patch and regression tests.
## Environment and verification
- Python 3.13.5, PyTorch 2.10.0+cpu, TorchMetrics 1.9.0, Linux, Gloo, two processes.
- The ten proposed local regression configurations give 3 failures / 7 passes
without the change and 10 passes with the change.
- The two-process reproduction succeeds after applying the same method change.
- The unguarded sequence is also present in the inspected `master` source at
`8d008de1660b18ba44fb1596f8a5e9e8361ba55c`. I have not run the complete test suite from that checkout.
Related history: #302 introduced `sync_context`; #339 added the explicit
synchronization-state logic.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with Metric.sync_context and unsync in the master source at the commit mentioned, and compare the synchronization-state logic introduced by #302 and #339. Run the two-process CPU reproduction and the local regression configurations. Done means the original compute() error still propagates, local state is restored, and a subsequent update/compute returns an MRR near 0.33333334 without a synchronization-state error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 65/100