Support pytorch_lightning with torch 2.6

Author: chezou

poetry.lock

@@ -88,7 +88,10 @@ torch = [
     {version = ">=1.6.0, <3.0.0", python = "<3.13", optional = true},
     {version = ">=2.6.0, <3.0.0", python = ">=3.13", optional = true},
 ]
-pytorch-lightning = {version = ">=1.6.0, <3.0.0", python = "<3.13", optional = true}
+pytorch-lightning = [
+    {version = ">=1.6.0, <3.0.0", python = "<3.13", optional = true},
+    {version = ">=2.5.1, <3.0.0", python = ">=3.13", optional = true},
+]
 
 ipywidgets = {version = ">=7.7,<8.2", optional = true}
 plotly = {version="^5.22.0", optional = true}

pyproject.toml

@@ -12,46 +12,26 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import sys
 import typing as tp
 
 import numpy as np
 import pandas as pd
 import pytest
-
-try:
-    from lightning_fabric import seed_everything
-except ImportError:
-    pass
-
-try:
-    import pytorch_lightning  # noqa  # pylint: disable=unused-import
-
-    filter_warnings_decorator = pytest.mark.filterwarnings(
-        "ignore::pytorch_lightning.utilities.warnings.PossibleUserWarning"
-    )
-except ImportError:
-
-    def filter_warnings_decorator(func):  # type: ignore
-        return func
-
+import pytorch_lightning  # noqa  # pylint: disable=unused-import
+from lightning_fabric import seed_everything
 
 from rectools.columns import Columns
 from rectools.dataset import Dataset
 from rectools.exceptions import NotFittedError
-
-try:
-    from rectools.models import DSSMModel
-    from rectools.models.nn.dssm import DSSM
-except ModuleNotFoundError:
-    pass
+from rectools.models import DSSMModel
+from rectools.models.nn.dssm import DSSM
 from rectools.models.vector import ImplicitRanker
 from tests.models.utils import assert_dumps_loads_do_not_change_model, assert_second_fit_refits_model
 
 from ..data import INTERACTIONS
 
-pytestmark = pytest.mark.skipif(
-    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
+filter_warnings_decorator = pytest.mark.filterwarnings(
+    "ignore::pytorch_lightning.utilities.warnings.PossibleUserWarning"
 )
 
 

tests/models/nn/test_dssm.py

@@ -12,18 +12,13 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import sys
 import typing as tp
 
 import numpy as np
 import pandas as pd
 import pytest
 import torch
-
-try:
-    from pytorch_lightning import seed_everything
-except ImportError:
-    pass
+from pytorch_lightning import seed_everything
 
 from rectools.columns import Columns
 from rectools.dataset import Dataset
@@ -38,10 +33,6 @@
 
 from ..data import DATASET, INTERACTIONS
 
-pytestmark = pytest.mark.skipif(
-    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
-)
-
 
 class TestIdEmbeddingsItemNet:
     def setup_method(self) -> None:

tests/models/nn/test_item_net.py

@@ -13,42 +13,25 @@
 #  limitations under the License.
 
 import os
-import sys
 import typing as tp
 from tempfile import NamedTemporaryFile
 
 import pandas as pd
 import pytest
 import torch
 from pytest import FixtureRequest
-
-try:
-    from pytorch_lightning import Trainer, seed_everything
-    from pytorch_lightning.loggers import CSVLogger
-
-except ImportError:
-    Trainer = object  # type: ignore
+from pytorch_lightning import Trainer, seed_everything
+from pytorch_lightning.loggers import CSVLogger
 
 from rectools import Columns
 from rectools.dataset import Dataset
-
-try:
-    from rectools.models import BERT4RecModel, SASRecModel, load_model
-    from rectools.models.nn.item_net import CatFeaturesItemNet, IdEmbeddingsItemNet
-    from rectools.models.nn.transformers.base import TransformerModelBase
-except ImportError:
-    TransformerModelBase = object  # type: ignore
+from rectools.models import BERT4RecModel, SASRecModel, load_model
+from rectools.models.nn.item_net import CatFeaturesItemNet, IdEmbeddingsItemNet
+from rectools.models.nn.transformers.base import TransformerModelBase
 from tests.models.data import INTERACTIONS
 from tests.models.utils import assert_save_load_do_not_change_model
 
-try:
-    from .utils import custom_trainer, custom_trainer_ckpt, custom_trainer_multiple_ckpt, leave_one_out_mask
-except NameError:
-    pass
-
-pytestmark = pytest.mark.skipif(
-    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
-)
+from .utils import custom_trainer, custom_trainer_ckpt, custom_trainer_multiple_ckpt, leave_one_out_mask
 
 
 class TestTransformerModelBase:

tests/models/nn/transformers/test_base.py

@@ -12,50 +12,36 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import sys
 import typing as tp
 from functools import partial
 
 import numpy as np
 import pandas as pd
 import pytest
 import torch
-
-try:
-    from pytorch_lightning import Trainer, seed_everything
-except ImportError:
-    Trainer = object  # type: ignore
+from pytorch_lightning import Trainer, seed_everything
 
 from rectools import ExternalIds
 from rectools.columns import Columns
 from rectools.dataset import Dataset
-
-try:
-    from rectools.models import BERT4RecModel
-    from rectools.models.nn.item_net import IdEmbeddingsItemNet, SumOfEmbeddingsConstructor
-    from rectools.models.nn.transformers.base import (
-        LearnableInversePositionalEncoding,
-        PreLNTransformerLayers,
-        TrainerCallable,
-        TransformerLightningModule,
-    )
-    from rectools.models.nn.transformers.bert4rec import MASKING_VALUE, BERT4RecDataPreparator, ValMaskCallable
-except ImportError:
-    TrainerCallable = object  # type: ignore
-    BERT4RecDataPreparator = object  # type: ignore
+from rectools.models import BERT4RecModel
+from rectools.models.nn.item_net import IdEmbeddingsItemNet, SumOfEmbeddingsConstructor
+from rectools.models.nn.transformers.base import (
+    LearnableInversePositionalEncoding,
+    PreLNTransformerLayers,
+    TrainerCallable,
+    TransformerLightningModule,
+)
+from rectools.models.nn.transformers.bert4rec import MASKING_VALUE, BERT4RecDataPreparator, ValMaskCallable
 from tests.models.data import DATASET
 from tests.models.utils import (
     assert_default_config_and_default_model_params_are_the_same,
     assert_second_fit_refits_model,
 )
 
-try:
-    from .utils import custom_trainer, leave_one_out_mask
-except NameError:
-    pass
+from .utils import custom_trainer, leave_one_out_mask
 
 
-@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13")
 class TestBERT4RecModel:
     def setup_method(self) -> None:
         self._seed_everything()
@@ -627,7 +613,6 @@ def _collate_fn_train(
         )
 
 
-@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13")
 class TestBERT4RecDataPreparator:
 
     def setup_method(self) -> None:
@@ -807,7 +792,6 @@ def test_get_dataloader_val(
             assert torch.equal(value, val_batch[key])
 
 
-@pytest.mark.skipif(sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13")
 class TestBERT4RecModelConfiguration:
     def setup_method(self) -> None:
         self._seed_everything()

tests/models/nn/transformers/test_bert4rec.py

@@ -14,52 +14,35 @@
 
 # pylint: disable=too-many-lines
 
-import sys
 import typing as tp
 from functools import partial
 
 import numpy as np
 import pandas as pd
 import pytest
 import torch
-
-try:
-    from pytorch_lightning import Trainer, seed_everything
-except ImportError:
-    Trainer = object  # type: ignore
+from pytorch_lightning import Trainer, seed_everything
 
 from rectools import ExternalIds
 from rectools.columns import Columns
 from rectools.dataset import Dataset, IdMap, Interactions
-
-try:
-    from rectools.models import SASRecModel
-    from rectools.models.nn.item_net import CatFeaturesItemNet, IdEmbeddingsItemNet, SumOfEmbeddingsConstructor
-    from rectools.models.nn.transformers.base import (
-        LearnableInversePositionalEncoding,
-        TrainerCallable,
-        TransformerLightningModule,
-        TransformerTorchBackbone,
-    )
-    from rectools.models.nn.transformers.sasrec import SASRecDataPreparator, SASRecTransformerLayers
-except ImportError:
-    TrainerCallable = object  # type: ignore
-    SASRecDataPreparator = object  # type: ignore
+from rectools.models import SASRecModel
+from rectools.models.nn.item_net import CatFeaturesItemNet, IdEmbeddingsItemNet, SumOfEmbeddingsConstructor
+from rectools.models.nn.transformers.base import (
+    LearnableInversePositionalEncoding,
+    TrainerCallable,
+    TransformerLightningModule,
+    TransformerTorchBackbone,
+)
+from rectools.models.nn.transformers.sasrec import SASRecDataPreparator, SASRecTransformerLayers
 from tests.models.data import DATASET
 from tests.models.utils import (
     assert_default_config_and_default_model_params_are_the_same,
     assert_second_fit_refits_model,
 )
 from tests.testing_utils import assert_id_map_equal, assert_interactions_set_equal
 
-try:
-    from .utils import custom_trainer, leave_one_out_mask
-except NameError:
-    pass
-
-pytestmark = pytest.mark.skipif(
-    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
-)
+from .utils import custom_trainer, leave_one_out_mask
 
 
 class TestSASRecModel:

tests/models/nn/transformers/test_sasrec.py

@@ -12,23 +12,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import sys
-
 import pandas as pd
-import pytest
-
-try:
-    from pytorch_lightning import Trainer
-    from pytorch_lightning.callbacks import ModelCheckpoint
-except ImportError:
-    pass
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
 
 from rectools import Columns
 
-pytestmark = pytest.mark.skipif(
-    sys.version_info >= (3, 13), reason="`pytorch_lightning` is not compatible with Python >= 3.13"
-)
-
 
 def leave_one_out_mask(interactions: pd.DataFrame) -> pd.Series:
     rank = (