Lightning-AI / Lightning-AI/torchmetrics

`Metric.sync_context` leaves the metric synchronized when `compute()` raises

Open
#3,486 1 comment 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.