fix version + yapf (#6999)

Author: Borda

CHANGELOG.md

@@ -11,6 +11,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed the order to call for world ranks & the `root_device` property in `TPUSpawnPlugin` ([#7074](https://github.com/PyTorchLightning/pytorch-lightning/pull/7074))
 - Fixed multi-gpu join for Horovod ([#6954](https://github.com/PyTorchLightning/pytorch-lightning/pull/6954))
+- Fixed parsing for pre-release package versions ([#6999](https://github.com/PyTorchLightning/pytorch-lightning/pull/6999))
+
 
 ## [1.2.8] - 2021-04-14
 
@@ -37,7 +39,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     * Remove hardcoding of local rank in accelerator connector ([#6878](https://github.com/PyTorchLightning/pytorch-lightning/pull/6878))
 
 
-
 ## [1.2.7] - 2021-04-06
 
 ### Fixed

pl_examples/basic_examples/dali_image_classifier.py

@@ -13,12 +13,12 @@
 # limitations under the License.
 from abc import ABC
 from argparse import ArgumentParser
-from distutils.version import LooseVersion
 from random import shuffle
 from warnings import warn
 
 import numpy as np
 import torch
+from packaging.version import Version
 from torch.nn import functional as F
 from torch.utils.data import random_split
 
@@ -44,7 +44,7 @@
     from nvidia.dali.pipeline import Pipeline
     from nvidia.dali.plugin.pytorch import DALIClassificationIterator
 
-    NEW_DALI_API = LooseVersion(dali_version) >= LooseVersion('0.28.0')
+    NEW_DALI_API = Version(dali_version) >= Version('0.28.0')
     if NEW_DALI_API:
         from nvidia.dali.plugin.base_iterator import LastBatchPolicy
 else:

pytorch_lightning/utilities/cloud_io.py

@@ -14,12 +14,12 @@
 
 import os
 import io
-from distutils.version import LooseVersion
 from pathlib import Path
 from typing import IO, Union
 
 import fsspec
 from fsspec.implementations.local import LocalFileSystem
+from packaging.version import Version
 
 import torch
 
@@ -71,7 +71,7 @@ def atomic_save(checkpoint, filepath: str):
     # Can't use the new zipfile serialization for 1.6.0 because there's a bug in
     # torch.hub.load_state_dict_from_url() that prevents it from loading the new files.
     # More details can be found here: https://github.com/pytorch/pytorch/issues/42239
-    if LooseVersion(torch.__version__).version[:3] == [1, 6, 0]:
+    if Version(torch.__version__).release[:3] == (1, 6, 0):
         torch.save(checkpoint, bytesbuffer, _use_new_zipfile_serialization=False)
     else:
         torch.save(checkpoint, bytesbuffer)

pytorch_lightning/utilities/imports.py

@@ -15,10 +15,10 @@
 import importlib
 import operator
 import platform
-from distutils.version import LooseVersion
 from importlib.util import find_spec
 
 import torch
+from packaging.version import Version
 from pkg_resources import DistributionNotFound
 
 
@@ -53,13 +53,11 @@ def _compare_version(package: str, op, version) -> bool:
     except (ModuleNotFoundError, DistributionNotFound):
         return False
     try:
-        pkg_version = LooseVersion(pkg.__version__)
-    except AttributeError:
-        return False
-    if not (hasattr(pkg_version, "vstring") and hasattr(pkg_version, "version")):
+        pkg_version = Version(pkg.__version__)
+    except TypeError:
         # this is mock by sphinx, so it shall return True ro generate all summaries
         return True
-    return op(pkg_version, LooseVersion(version))
+    return op(pkg_version, Version(version))
 
 
 _IS_WINDOWS = platform.system() == "Windows"

requirements.txt

@@ -7,5 +7,5 @@ future>=0.17.1  # required for builtins in setup.py
 PyYAML>=5.1, !=5.4.*  # OmegaConf requirement >=5.1
 tqdm>=4.41.0
 fsspec[http]>=0.8.1
-tensorboard>=2.2.0
-torchmetrics>=0.2.0
+tensorboard>=2.2.0, !=2.5.0  # 2.5.0 GPU CI error: 'Couldn't build proto file into descriptor pool!'
+torchmetrics==0.2.0

requirements/extra.txt

@@ -8,4 +8,4 @@ torchtext>=0.5
 onnxruntime>=1.3.0
 hydra-core>=1.0
 # todo: when switch to standard package stream, drop `fairscale` from hard mocked docs libs
-https://github.com/PyTorchLightning/fairscale/archive/pl_1.2.0.zip
+fairscale @ https://github.com/PyTorchLightning/fairscale/archive/pl_1.2.0.zip

tests/helpers/runif.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 import os
 import sys
-from distutils.version import LooseVersion
 from typing import Optional
 
 import pytest
 import torch
+from packaging.version import Version
 from pkg_resources import get_distribution
 
 from pytorch_lightning.utilities import (
@@ -100,18 +100,18 @@ def __new__(
             reasons.append(f"GPUs>={min_gpus}")
 
         if min_torch:
-            torch_version = LooseVersion(get_distribution("torch").version)
-            conditions.append(torch_version < LooseVersion(min_torch))
+            torch_version = get_distribution("torch").version
+            conditions.append(Version(torch_version) < Version(min_torch))
             reasons.append(f"torch>={min_torch}")
 
         if max_torch:
-            torch_version = LooseVersion(get_distribution("torch").version)
-            conditions.append(torch_version >= LooseVersion(max_torch))
+            torch_version = get_distribution("torch").version
+            conditions.append(Version(torch_version) >= Version(max_torch))
             reasons.append(f"torch<{max_torch}")
 
         if min_python:
             py_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
-            conditions.append(py_version < LooseVersion(min_python))
+            conditions.append(Version(py_version) < Version(min_python))
             reasons.append(f"python>={min_python}")
 
         if quantization:

tests/loggers/test_tensorboard.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 import os
 from argparse import Namespace
-from distutils.version import LooseVersion
 from unittest import mock
 
 import pytest
 import torch
 import yaml
 from omegaconf import OmegaConf
+from packaging.version import Version
 from tensorboard.backend.event_processing.event_accumulator import EventAccumulator
 
 from pytorch_lightning import Trainer
@@ -65,7 +65,7 @@ def __init__(self, b1=0.5, b2=0.999):
 
     data_pt_1_5 = b'\x12\x1b"\x04\n\x02b1"\x04\n\x02b2*\r\n\x0b\x12\thp_metric'
     data_pt_1_6 = b'\x12\x1f"\x06\n\x02b1 \x03"\x06\n\x02b2 \x03*\r\n\x0b\x12\thp_metric'
-    hparams_data = data_pt_1_6 if LooseVersion(torch.__version__) >= LooseVersion("1.6.0") else data_pt_1_5
+    hparams_data = data_pt_1_6 if Version(torch.__version__) >= Version("1.6.0") else data_pt_1_5
 
     assert event_acc.summary_metadata['_hparams_/experiment'].plugin_data.plugin_name == 'hparams'
     assert event_acc.summary_metadata['_hparams_/experiment'].plugin_data.content == hparams_data