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

[hse3]

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:

python -m torch.distributed.run --standalone --nproc_per_node=2 repro_compute.py

"""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:

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.

[github-actions]

Hi! Thanks for your contribution! Great first issue!