Repository: aimhubio/aim

Issue:
Handling for aliases of PyTorch Lightning
## 🐛 Bug

Starting from PyTorch Lightning 1.8, this package has been renamed to `Lightning`. While both `lightning` and `pytorch_lightning` are in use as a package (`lightning` is a superset package for `pytorch_lightning` now).

This has caused a problem that `aim`'s code only support for PyTorch Lightning when `pytorch_lightning` is installed. When `lightning` is installed for some other users/developers, they will find that `aim` will wrongly suggest that they have not installed `pytorch_lightning` package.

Since parallel distribution for both `lightning` and `pytorch_lightning` would [stays for a long time](https://github.com/Lightning-AI/lightning/discussions/17095). It may be necessary to support both packages in `aim`'s code.

### To reproduce

Use `pip install lightning` instead of `pip install pytorch_lightning` to install `PyTorch Lightning`.
### Expected behavior

pytorch lightning addons should work for both those who installed `lightning` package and `pytorch-lightning` package.

### Environment

- Aim Version (e.g., 3.0.1): latest
- Python version: 3.10.1
- pip version
- OS (e.g., Linux): Linux
- PyTorch Lightning >= 1.8.0

### Additional context

None



Discussion:
I'll try to submit a PR to fix this issue in `src/aim/pytorch_lightning.py`
A corresponding fixing PR has been submitted #2747 

Resolving patch:
diff --git a/aim/sdk/adapters/pytorch_lightning.py b/aim/sdk/adapters/pytorch_lightning.py
--- a/aim/sdk/adapters/pytorch_lightning.py
+++ b/aim/sdk/adapters/pytorch_lightning.py
@@ -1,10 +1,19 @@
 import os
+import importlib.util
 from typing import Any, Dict, Optional, Union
 from argparse import Namespace
 
 import packaging.version
 
-try:
+if importlib.util.find_spec("lightning"):
+    import lightning.pytorch as pl
+
+    from lightning.pytorch.loggers.logger import (
+        Logger, rank_zero_experiment
+    )
+
+    from lightning.pytorch.utilities import rank_zero_only
+elif importlib.util.find_spec("pytorch_lightning"):
     import pytorch_lightning as pl
 
     if packaging.version.parse(pl.__version__) < packaging.version.parse("1.7"):
@@ -19,10 +28,11 @@
         )
 
     from pytorch_lightning.utilities import rank_zero_only
-except ImportError:
+else:
     raise RuntimeError(
         'This contrib module requires PyTorch Lightning to be installed. '
         'Please install it with command: \n pip install pytorch-lightning'
+        'or \n pip install lightning'
     )
 
 from aim.sdk.run import Run
diff --git a/examples/pytorch_lightning_track.py b/examples/pytorch_lightning_track.py
--- a/examples/pytorch_lightning_track.py
+++ b/examples/pytorch_lightning_track.py
@@ -1,9 +1,19 @@
+import importlib.util
 from aim.pytorch_lightning import AimLogger
 
 from argparse import ArgumentParser
 
 import torch
-import pytorch_lightning as pl
+if importlib.util.find_spec("lightning"):
+    import lightning.pytorch as pl
+elif importlib.util.find_spec("pytorch_lightning"): # noqa F401
+    import pytorch_lightning as pl
+else:
+    raise RuntimeError(
+        'This contrib module requires PyTorch Lightning to be installed. '
+        'Please install it with command: \n pip install pytorch-lightning \n'
+        'or \n pip install lightning'
+    )
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, random_split
 


Test patch: