Merge branch 'master' into device-enhance

zhiyuan1i · web-flow · commit dee68b19cc5f · 2024-12-25T17:11:08.000+08:00
diff --git a/.actions/assistant.py b/.actions/assistant.py
@@ -481,23 +481,27 @@ def convert_version2nightly(ver_file: str = "src/version.info") -> None:
         with open(ver_file, "w") as fo:
             fo.write(version + os.linesep)
 
+    @staticmethod
+    def generate_docker_tags(
+        release_version: str,
+        python_version: str,
+        torch_version: str,
+        cuda_version: str,
+        docker_project: str = "pytorchlightning/pytorch_lightning",
+        add_latest: bool = False,
+    ) -> None:
+        """Generate docker tags for the given versions."""
+        tags = [f"latest-py{python_version}-torch{torch_version}-cuda{cuda_version}"]
+        if release_version:
+            tags += [f"{release_version}-py{python_version}-torch{torch_version}-cuda{cuda_version}"]
+        if add_latest:
+            tags += ["latest"]
 
-if __name__ == "__main__":
-    import sys
-
-    import jsonargparse
-    from jsonargparse import ArgumentParser
-
-    def patch_jsonargparse_python_3_12_8():
-        if sys.version_info < (3, 12, 8):
-            return
-
-        def _parse_known_args_patch(self: ArgumentParser, args: Any = None, namespace: Any = None) -> tuple[Any, Any]:
-            namespace, args = super(ArgumentParser, self)._parse_known_args(args, namespace, intermixed=False)  # type: ignore
-            return namespace, args
+        tags = [f"{docker_project}:{tag}" for tag in tags]
+        print(",".join(tags))
 
-        setattr(ArgumentParser, "_parse_known_args", _parse_known_args_patch)
 
-    patch_jsonargparse_python_3_12_8()  # Required until fix https://github.com/omni-us/jsonargparse/issues/641
+if __name__ == "__main__":
+    import jsonargparse
 
     jsonargparse.CLI(AssistantCLI, as_positional=False)
diff --git a/.actions/requirements.txt b/.actions/requirements.txt
@@ -1,3 +1,3 @@
-jsonargparse >=4.16.0, <4.28.0
+jsonargparse >=4.16.0, <=4.35.0
 requests
 packaging
diff --git a/.azure/gpu-benchmarks.yml b/.azure/gpu-benchmarks.yml
@@ -77,7 +77,7 @@ jobs:
 
       - bash: |
           pip install -e .[dev] --find-links ${TORCH_URL}
-          pip install setuptools==75.6.0
+          pip install setuptools==75.6.0 jsonargparse==4.35.0
         env:
           FREEZE_REQUIREMENTS: "1"
         displayName: "Install package"
diff --git a/.azure/gpu-tests-fabric.yml b/.azure/gpu-tests-fabric.yml
@@ -107,7 +107,7 @@ jobs:
       - bash: |
           extra=$(python -c "print({'lightning': 'fabric-'}.get('$(PACKAGE_NAME)', ''))")
           pip install -e ".[${extra}dev]" pytest-timeout -U --find-links="${TORCH_URL}" --find-links="${TORCHVISION_URL}"
-          pip install setuptools==75.6.0
+          pip install setuptools==75.6.0 jsonargparse==4.35.0
         displayName: "Install package & dependencies"
 
       - bash: |
diff --git a/.azure/gpu-tests-pytorch.yml b/.azure/gpu-tests-pytorch.yml
@@ -111,7 +111,7 @@ jobs:
       - bash: |
           extra=$(python -c "print({'lightning': 'pytorch-'}.get('$(PACKAGE_NAME)', ''))")
           pip install -e ".[${extra}dev]" pytest-timeout -U --find-links="${TORCH_URL}" --find-links="${TORCHVISION_URL}"
-          pip install setuptools==75.6.0
+          pip install setuptools==75.6.0 jsonargparse==4.35.0
         displayName: "Install package & dependencies"
 
       - bash: pip uninstall -y lightning
diff --git a/.github/checkgroup.yml b/.github/checkgroup.yml
@@ -145,17 +145,17 @@ subprojects:
       - "!*.md"
       - "!**/*.md"
     checks:
-      - "build-cuda (3.10, 2.1.2, 12.1.0)"
-      - "build-cuda (3.11, 2.2.2, 12.1.0)"
-      - "build-cuda (3.11, 2.3.1, 12.1.0)"
-      - "build-cuda (3.11, 2.4.1, 12.1.0)"
-      - "build-cuda (3.12, 2.5.1, 12.1.0)"
+      - "build-cuda (3.10, 2.1.2, 12.1.1)"
+      - "build-cuda (3.11, 2.2.2, 12.1.1)"
+      - "build-cuda (3.11, 2.3.1, 12.1.1)"
+      - "build-cuda (3.11, 2.4.1, 12.1.1)"
+      - "build-cuda (3.12, 2.5.1, 12.1.1)"
       #- "build-NGC"
-      - "build-pl (3.10, 2.1, 12.1.0)"
-      - "build-pl (3.11, 2.2, 12.1.0)"
-      - "build-pl (3.11, 2.3, 12.1.0)"
-      - "build-pl (3.11, 2.4, 12.1.0)"
-      - "build-pl (3.12, 2.5, 12.1.0)"
+      - "build-pl (3.10, 2.1, 12.1.1)"
+      - "build-pl (3.11, 2.2, 12.1.1)"
+      - "build-pl (3.11, 2.3, 12.1.1)"
+      - "build-pl (3.11, 2.4, 12.1.1)"
+      - "build-pl (3.12, 2.5, 12.1.1, true)"
 
   # SECTION: lightning_fabric
 
diff --git a/.github/workflows/ci-tests-fabric.yml b/.github/workflows/ci-tests-fabric.yml
@@ -49,16 +49,16 @@ jobs:
           - { os: "macOS-14", pkg-name: "lightning", python-version: "3.11", pytorch-version: "2.3" }
           - { os: "ubuntu-20.04", pkg-name: "lightning", python-version: "3.11", pytorch-version: "2.3" }
           - { os: "windows-2022", pkg-name: "lightning", python-version: "3.11", pytorch-version: "2.3" }
-          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.4.1" }
-          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.4.1" }
-          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.4.1" }
-          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.5.1" }
+          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.4.1" }
+          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.4.1" }
+          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.4.1" }
+          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.5.1" }
           # only run PyTorch latest with Python latest, use Fabric scope to limit dependency issues
-          - { os: "macOS-14", pkg-name: "fabric", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "ubuntu-22.04", pkg-name: "fabric", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "windows-2022", pkg-name: "fabric", python-version: "3.12", pytorch-version: "2.5.1" }
+          - { os: "macOS-14", pkg-name: "fabric", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "ubuntu-22.04", pkg-name: "fabric", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "windows-2022", pkg-name: "fabric", python-version: "3.12.7", pytorch-version: "2.5.1" }
           # "oldest" versions tests, only on minimum Python
           - { os: "macOS-14", pkg-name: "lightning", python-version: "3.9", pytorch-version: "2.1", requires: "oldest" }
           - {
diff --git a/.github/workflows/ci-tests-pytorch.yml b/.github/workflows/ci-tests-pytorch.yml
@@ -53,16 +53,16 @@ jobs:
           - { os: "macOS-14", pkg-name: "lightning", python-version: "3.11", pytorch-version: "2.3" }
           - { os: "ubuntu-20.04", pkg-name: "lightning", python-version: "3.11", pytorch-version: "2.3" }
           - { os: "windows-2022", pkg-name: "lightning", python-version: "3.11", pytorch-version: "2.3" }
-          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.4.1" }
-          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.4.1" }
-          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.4.1" }
-          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12", pytorch-version: "2.5.1" }
+          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.4.1" }
+          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.4.1" }
+          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.4.1" }
+          - { os: "macOS-14", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "ubuntu-22.04", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "windows-2022", pkg-name: "lightning", python-version: "3.12.7", pytorch-version: "2.5.1" }
           # only run PyTorch latest with Python latest, use PyTorch scope to limit dependency issues
-          - { os: "macOS-14", pkg-name: "pytorch", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "ubuntu-22.04", pkg-name: "pytorch", python-version: "3.12", pytorch-version: "2.5.1" }
-          - { os: "windows-2022", pkg-name: "pytorch", python-version: "3.12", pytorch-version: "2.5.1" }
+          - { os: "macOS-14", pkg-name: "pytorch", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "ubuntu-22.04", pkg-name: "pytorch", python-version: "3.12.7", pytorch-version: "2.5.1" }
+          - { os: "windows-2022", pkg-name: "pytorch", python-version: "3.12.7", pytorch-version: "2.5.1" }
           # "oldest" versions tests, only on minimum Python
           - { os: "macOS-14", pkg-name: "lightning", python-version: "3.9", pytorch-version: "2.1", requires: "oldest" }
           - {
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -43,11 +43,11 @@ jobs:
         include:
           # We only release one docker image per PyTorch version.
           # Make sure the matrix here matches the one below.
-          - { python_version: "3.10", pytorch_version: "2.1", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.2", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.3", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.4", cuda_version: "12.1.0" }
-          - { python_version: "3.12", pytorch_version: "2.5", cuda_version: "12.1.0" }
+          - { python_version: "3.10", pytorch_version: "2.1", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.2", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.3", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.4", cuda_version: "12.1.1" }
+          - { python_version: "3.12", pytorch_version: "2.5", cuda_version: "12.1.1", latest: "true" }
     steps:
       - uses: actions/checkout@v4
         with:
@@ -65,23 +65,14 @@ jobs:
         run: echo "RELEASE_VERSION=$(echo ${GITHUB_REF##*/})" >> $GITHUB_ENV
       - name: Set tags
         run: |
-          import os
-
-          repo = "pytorchlightning/pytorch_lightning"
-          ver = os.getenv('RELEASE_VERSION')
-          py_ver = "${{ matrix.python_version }}"
-          pt_ver = "${{ matrix.pytorch_version }}"
-          cuda_ver = "${{ matrix.cuda_version }}"
-          tags = [f"latest-py{py_ver}-torch{pt_ver}-cuda{cuda_ver}"]
-          if ver:
-            tags += [f"{ver}-py{py_ver}-torch{pt_ver}-cuda{cuda_ver}"]
-          if py_ver == '3.11' and pt_ver == '2.3' and cuda_ver == '12.1.0':
-            tags += ["latest"]
-
-          tags = [f"{repo}:{tag}" for tag in tags]
-          with open(os.getenv('GITHUB_ENV'), "a") as gh_env:
-              gh_env.write("DOCKER_TAGS=" + ",".join(tags))
-        shell: python
+          pip install -q -r .actions/requirements.txt
+          tags=$(python .actions/assistant.py generate_docker_tags \
+            --release_version="${{ env.RELEASE_VERSION }}" \
+            --python_version="${{ matrix.python_version }}" \
+            --torch_version="${{ matrix.pytorch_version }}" \
+            --cuda_version="${{ matrix.cuda_version }}" \
+            --add_latest="${{ matrix.latest || 'false' }}")
+          echo "DOCKER_TAGS=$tags" >> $GITHUB_ENV
 
       - uses: docker/build-push-action@v6
         with:
@@ -104,11 +95,11 @@ jobs:
         include:
           # These are the base images for PL release docker images.
           # Make sure the matrix here matches the one above.
-          - { python_version: "3.10", pytorch_version: "2.1.2", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.2.2", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.3.1", cuda_version: "12.1.0" }
-          - { python_version: "3.11", pytorch_version: "2.4.1", cuda_version: "12.1.0" }
-          - { python_version: "3.12", pytorch_version: "2.5.1", cuda_version: "12.1.0" }
+          - { python_version: "3.10", pytorch_version: "2.1.2", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.2.2", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.3.1", cuda_version: "12.1.1" }
+          - { python_version: "3.11", pytorch_version: "2.4.1", cuda_version: "12.1.1" }
+          - { python_version: "3.12", pytorch_version: "2.5.1", cuda_version: "12.1.1" }
     steps:
       - uses: actions/checkout@v4
       - uses: docker/setup-buildx-action@v3
diff --git a/docs/source-pytorch/conf.py b/docs/source-pytorch/conf.py
@@ -464,8 +464,8 @@ def _load_py_module(name: str, location: str) -> ModuleType:
     ("py:func", "lightning.pytorch.utilities.rank_zero.rank_zero_only"),
     ("py:class", "lightning.pytorch.utilities.types.LRSchedulerConfig"),
     ("py:class", "lightning.pytorch.utilities.types.LRSchedulerConfigType"),
-    ("py:class", "lightning.pytorch.utilities.types.OptimizerConfigType"),
-    ("py:class", "lightning.pytorch.utilities.types.OptimizerLRSchedulerConfigType"),
+    ("py:class", "lightning.pytorch.utilities.types.OptimizerConfig"),
+    ("py:class", "lightning.pytorch.utilities.types.OptimizerLRSchedulerConfig"),
     ("py:class", "lightning_habana.pytorch.plugins.precision.HPUPrecisionPlugin"),
     ("py:class", "lightning_habana.pytorch.strategies.HPUDDPStrategy"),
     ("py:class", "lightning_habana.pytorch.strategies.HPUParallelStrategy"),
diff --git a/requirements/pytorch/extra.txt b/requirements/pytorch/extra.txt
@@ -5,7 +5,7 @@
 matplotlib>3.1, <3.9.0
 omegaconf >=2.2.3, <2.4.0
 hydra-core >=1.2.0, <1.4.0
-jsonargparse[signatures] >=4.27.7, <4.28.0
+jsonargparse[signatures] >=4.27.7, <=4.35.0
 rich >=12.3.0, <13.6.0
 tensorboardX >=2.2, <2.7.0  # min version is set by torch.onnx missing attribute
 bitsandbytes >=0.44.0,<0.44.2; sys_platform == 'linux' or sys_platform == 'win32'
diff --git a/src/lightning/pytorch/cli.py b/src/lightning/pytorch/cli.py
@@ -37,18 +37,6 @@
 
 _JSONARGPARSE_SIGNATURES_AVAILABLE = RequirementCache("jsonargparse[signatures]>=4.27.7")
 
-
-def patch_jsonargparse_python_3_12_8() -> None:
-    if sys.version_info < (3, 12, 8):
-        return
-
-    def _parse_known_args_patch(self: ArgumentParser, args: Any = None, namespace: Any = None) -> tuple[Any, Any]:
-        namespace, args = super(ArgumentParser, self)._parse_known_args(args, namespace, intermixed=False)  # type: ignore
-        return namespace, args
-
-    setattr(ArgumentParser, "_parse_known_args", _parse_known_args_patch)
-
-
 if _JSONARGPARSE_SIGNATURES_AVAILABLE:
     import docstring_parser
     from jsonargparse import (
@@ -60,8 +48,6 @@ def _parse_known_args_patch(self: ArgumentParser, args: Any = None, namespace: A
         set_config_read_mode,
     )
 
-    patch_jsonargparse_python_3_12_8()  # Required until fix https://github.com/omni-us/jsonargparse/issues/641
-
     register_unresolvable_import_paths(torch)  # Required until fix https://github.com/pytorch/pytorch/issues/74483
     set_config_read_mode(fsspec_enabled=True)
 else:
diff --git a/src/lightning/pytorch/utilities/data.py b/src/lightning/pytorch/utilities/data.py
@@ -351,7 +351,11 @@ def _is_dataloader_shuffled(dataloader: object) -> bool:
         # shuffling is enabled via a sampler. No sampler, no shuffling
         return False
     batch_sampler = dataloader.batch_sampler
-    sampler = batch_sampler.sampler if batch_sampler is not None else dataloader.sampler
+    if batch_sampler is not None:
+        # custom batch samplers may not have an internal .sampler
+        sampler = batch_sampler.sampler if hasattr(batch_sampler, "sampler") else batch_sampler
+    else:
+        sampler = dataloader.sampler
     if isinstance(sampler, SequentialSampler):
         return False
     return isinstance(sampler, RandomSampler)
diff --git a/src/lightning/pytorch/utilities/types.py b/src/lightning/pytorch/utilities/types.py
@@ -104,11 +104,11 @@ class LRSchedulerConfigType(TypedDict, total=False):
     strict: bool
 
 
-class OptimizerConfigType(TypedDict):
+class OptimizerConfig(TypedDict):
     optimizer: Optimizer
 
 
-class OptimizerLRSchedulerConfigType(TypedDict):
+class OptimizerLRSchedulerConfig(TypedDict):
     optimizer: Optimizer
     lr_scheduler: Union[LRSchedulerTypeUnion, LRSchedulerConfigType]
     monitor: NotRequired[str]
@@ -119,10 +119,10 @@ class OptimizerLRSchedulerConfigType(TypedDict):
         Optimizer,
         Sequence[Optimizer],
         tuple[Sequence[Optimizer], Sequence[Union[LRSchedulerTypeUnion, LRSchedulerConfig]]],
-        OptimizerConfigType,
-        OptimizerLRSchedulerConfigType,
-        Sequence[OptimizerConfigType],
-        Sequence[OptimizerLRSchedulerConfigType],
+        OptimizerConfig,
+        OptimizerLRSchedulerConfig,
+        Sequence[OptimizerConfig],
+        Sequence[OptimizerLRSchedulerConfig],
     ]
 ]
 
diff --git a/src/version.info b/src/version.info
@@ -1 +1 @@
-2.5.0
+2.5.0.post0
diff --git a/tests/parity_fabric/test_parity_ddp.py b/tests/parity_fabric/test_parity_ddp.py
@@ -162,8 +162,5 @@ def run_parity_test(accelerator: str = "cpu", devices: int = 2, tolerance: float
 
 if __name__ == "__main__":
     from jsonargparse import CLI
-    from lightning.pytorch.cli import patch_jsonargparse_python_3_12_8
-
-    patch_jsonargparse_python_3_12_8()  # Required until fix https://github.com/omni-us/jsonargparse/issues/641
 
     CLI(run_parity_test)
diff --git a/tests/tests_pytorch/checkpointing/test_model_checkpoint.py b/tests/tests_pytorch/checkpointing/test_model_checkpoint.py
@@ -29,10 +29,10 @@
 import pytest
 import torch
 import yaml
+from jsonargparse import ArgumentParser
 from lightning.fabric.utilities.cloud_io import _load as pl_load
 from lightning.pytorch import Trainer, seed_everything
 from lightning.pytorch.callbacks import ModelCheckpoint
-from lightning.pytorch.cli import LightningArgumentParser as ArgumentParser
 from lightning.pytorch.demos.boring_classes import BoringModel
 from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
diff --git a/tests/tests_pytorch/test_cli.py b/tests/tests_pytorch/test_cli.py
@@ -754,6 +754,7 @@ def add_arguments_to_parser(self, parser):
 
 
 @_xfail_python_ge_3_11_9
+@RunIf(min_torch="2.2")
 @pytest.mark.parametrize("use_generic_base_class", [False, True])
 def test_lightning_cli_optimizers_and_lr_scheduler_with_link_to(use_generic_base_class):
     class MyLightningCLI(LightningCLI):
@@ -801,6 +802,7 @@ def __init__(self, optim1: dict, optim2: dict, scheduler: dict):
 
 
 @_xfail_python_ge_3_11_9
+@RunIf(min_torch="2.2")
 def test_lightning_cli_optimizers_and_lr_scheduler_with_callable_type():
     class TestModel(BoringModel):
         def __init__(
