Repository: ad12/meddlr (language: python)

Issue:
Invalid Version when Using Self-Compiled/Beta PyTorch
I encountered an exception when running `import meddlr.ops as oF` in an NGC PyTorch container.

The exception says,

```python
>>> import meddlr.ops as oF
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/workspace/meddlr/meddlr/__init__.py", line 3, in <module>
    from meddlr.engine.model_zoo import get_model_from_zoo
  File "/workspace/meddlr/meddlr/engine/__init__.py", line 10, in <module>
    from .hooks import *  # noqa
  File "/workspace/meddlr/meddlr/engine/hooks.py", line 19, in <module>
    from meddlr.evaluation.testing import flatten_results_dict
  File "/workspace/meddlr/meddlr/evaluation/__init__.py", line 2, in <module>
    from .recon_evaluation import ReconEvaluator  # noqa: F401
  File "/workspace/meddlr/meddlr/evaluation/recon_evaluation.py", line 16, in <module>
    from meddlr.data.transforms.transform import build_normalizer
  File "/workspace/meddlr/meddlr/data/__init__.py", line 1, in <module>
    from meddlr.data import build, catalog, collate, data_utils, slice_dataset
  File "/workspace/meddlr/meddlr/data/build.py", line 13, in <module>
    from meddlr.data.slice_dataset import SliceData
  File "/workspace/meddlr/meddlr/data/slice_dataset.py", line 10, in <module>
    from meddlr.data.transforms.transform import DataTransform
  File "/workspace/meddlr/meddlr/data/transforms/transform.py", line 9, in <module>
    from meddlr.forward import SenseModel
  File "/workspace/meddlr/meddlr/forward/__init__.py", line 1, in <module>
    from meddlr.forward import mri
  File "/workspace/meddlr/meddlr/forward/mri.py", line 4, in <module>
    import meddlr.ops as oF
  File "/workspace/meddlr/meddlr/ops/__init__.py", line 1, in <module>
    from meddlr.ops import categorical, complex, fft, utils  # noqa: F401
  File "/workspace/meddlr/meddlr/ops/fft.py", line 10, in <module>
    if env.pt_version() >= [1, 6]:
  File "/workspace/meddlr/meddlr/utils/env.py", line 314, in pt_version
    version = [dtype(x) for x in version]
  File "/workspace/meddlr/meddlr/utils/env.py", line 314, in <listcomp>
    version = [dtype(x) for x in version]
ValueError: invalid literal for int() with base 10: '0a0'
```

I checked the related lines, the code assumes the PyTorch is using stable versions and installed from binary packs.

https://github.com/ad12/meddlr/blob/b2f055b482c534a2a6eb161c19e6abb6ea94c196/meddlr/utils/env.py#L314

However, the PyTorch in the NGC container is self-complied, and the version string contains a commit ID. It looks like this,

```python
>>> import torch
>>> torch.__version__
'1.13.0a0+08820cb'
```

I suspect that's why it triggered the exception.

Would you please kindly modify the function to accept the version strings like this?


Discussion / hints:

Fix patch:
diff --git a/meddlr/utils/env.py b/meddlr/utils/env.py
index 8ceb152b..f2ecc1e8 100644
--- a/meddlr/utils/env.py
+++ b/meddlr/utils/env.py
@@ -10,7 +10,7 @@
 import warnings
 from datetime import datetime
 from importlib import util
-from typing import List
+from typing import Sequence, Union
 
 import numpy as np
 import torch
@@ -25,6 +25,38 @@
 _SUPPORTED_PACKAGES = {}
 
 
+class Version(version.Version):
+    """
+    An extension of packaging.version.Version that supports
+    comparisons with strings and lists.
+    """
+
+    def _format_version(self, other: Union[str, Sequence[Union[str, int]]]):
+        if isinstance(other, str):
+            return version.Version(other)
+        if isinstance(other, (list, tuple)):
+            return version.Version(".".join(map(str, other)))
+        return other
+
+    def __eq__(self, other):
+        return super().__eq__(self._format_version(other))
+
+    def __lt__(self, other):
+        return super().__lt__(self._format_version(other))
+
+    def __le__(self, other):
+        return super().__le__(self._format_version(other))
+
+    def __gt__(self, other):
+        return super().__gt__(self._format_version(other))
+
+    def __ge__(self, other):
+        return super().__ge__(self._format_version(other))
+
+    def __ne__(self, other):
+        return super().__ne__(self._format_version(other))
+
+
 def seed_all_rng(seed=None):
     """
     Set the random seed for the RNG in torch, numpy and python.
@@ -307,12 +339,9 @@ def supports_cupy():
     return package_available("cupy")
 
 
-def pt_version(dtype=int) -> List:
-    version = list(_PT_VERSION.split("."))
-    version[-1] = version[-1].split("+")[0]
-    if not issubclass(dtype, str):
-        version = [dtype(x) for x in version]
-    return version
+def pt_version() -> Version:
+    """Returns the PyTorch version."""
+    return Version(_PT_VERSION)
 
 
 def supports_cplx_tensor() -> bool:


Test patch:
diff --git a/tests/utils/test_env.py b/tests/utils/test_env.py
index 74b4260c..c487419e 100644
--- a/tests/utils/test_env.py
+++ b/tests/utils/test_env.py
@@ -1,6 +1,9 @@
 import os
 import unittest
 
+import pytest
+from packaging import version
+
 from meddlr.utils import env
 
 
@@ -92,5 +95,33 @@ def test_is_package_installed():
     assert env.is_package_installed(f"numpy=={numpy_version}")
 
 
+@pytest.mark.parametrize("comp_type", ["str", "int_list", "version"])
+def test_version(comp_type):
+    """Test that we can compare with list and strings."""
+
+    def _format_value(val):
+        if comp_type == "str":
+            return val
+        elif comp_type == "int_list":
+            return [int(x) for x in val.split(".")]
+        elif comp_type == "version":
+            return version.Version(val)
+        else:
+            raise ValueError(f"Unknown comp_type '{comp_type}'.")
+
+    curr_version = env.Version("1.0.0")
+    assert curr_version == _format_value("1.0.0")
+    assert curr_version > _format_value("0.9.9")
+    assert curr_version >= _format_value("0.9.9")
+    assert curr_version < _format_value("1.1.1")
+    assert curr_version <= _format_value("1.1.1")
+    assert curr_version != _format_value("1.1.1")
+
+
+def test_custom_pt_version_parsing():
+    """Test that custom builds of PyTorch can be parsed and compared."""
+    assert env.Version("1.13.0a0+08820cb") >= [1, 10]
+
+
 if __name__ == "__main__":
     unittest.main()