Repository: Lightning-AI/torchmetrics

Issue:
MetricCollection should enforce order when passed a dict
## 🐛 Bug

Not a clear bug yet, but just thinking about distributed metric computation, in here: https://github.com/PyTorchLightning/metrics/blob/53d570158a503497351ae45ec895ca44a0546068/torchmetrics/collections.py#L81
we should make sure to sort the names before the insert so that we insert things in the same order (ModuleDict is already OrderedDict otherwise). If we don't we will get deadlocks when doing distributed metric updates.

Additionally, we might want to enforce sorting when passed list/tuple, but that might be more on the user end.


### To Reproduce

On each of the workers, pass dictionary with same metrics but in different order, try compute and observe deadlock.




Discussion:
I agree that this should make the it more robust.
I think the only reason we have not seen this as a problem before, is that you explicit need to construct it in a different order on each process.
@maximsch2 after thinking a bit more about this, is it just an artificially problem?
You need to explicit create a different order dependent on the rank of the process for this to be a trouble.
I think it's not super easy issue to trigger with simple code, but if you have a framework that wraps your code and constructs metric collections, then things can go weird. Given that it's pretty simple problem to solve, I think it's worth it to be defensive here. Let me just send a PR actually.

Resolving patch:
diff --git a/torchmetrics/collections.py b/torchmetrics/collections.py
--- a/torchmetrics/collections.py
+++ b/torchmetrics/collections.py
@@ -35,6 +35,7 @@ class MetricCollection(nn.ModuleDict):
             * dict: if metrics are passed in as a dict, will use each key in the
               dict as key for output dict. Use this format if you want to chain
               together multiple of the same metric with different parameters.
+              Note that the keys in the output dict will be sorted alphabetically.
 
         prefix: a string to append in front of the keys of the output dict
 
@@ -78,7 +79,9 @@ def __init__(
         super().__init__()
         if isinstance(metrics, dict):
             # Check all values are metrics
-            for name, metric in metrics.items():
+            # Make sure that metrics are added in deterministic order
+            for name in sorted(metrics.keys()):
+                metric = metrics[name]
                 if not isinstance(metric, Metric):
                     raise ValueError(
                         f"Value {metric} belonging to key {name}"


Test patch:
diff --git a/tests/bases/test_collections.py b/tests/bases/test_collections.py
--- a/tests/bases/test_collections.py
+++ b/tests/bases/test_collections.py
@@ -156,3 +156,12 @@ def test_metric_collection_prefix_arg(tmpdir):
     out = new_metric_collection(5)
     for name in names:
         assert f"new_prefix_{name}" in out, 'prefix argument not working as intended with clone method'
+
+
+def test_metric_collection_same_order():
+    m1 = DummyMetricSum()
+    m2 = DummyMetricDiff()
+    col1 = MetricCollection({"a": m1, "b": m2})
+    col2 = MetricCollection({"b": m2, "a": m1})
+    for k1, k2 in zip(col1.keys(), col2.keys()):
+        assert k1 == k2