contextlib.AbstractContextManager

Author: Borda

src/lightning/fabric/fabric.py

@@ -14,13 +14,12 @@
 import inspect
 import os
 from collections.abc import Generator, Mapping, Sequence
-from contextlib import contextmanager, nullcontext
+from contextlib import AbstractContextManager, contextmanager, nullcontext
 from functools import partial
 from pathlib import Path
 from typing import (
     Any,
     Callable,
-    ContextManager,
     Optional,
     Union,
     cast,
@@ -484,7 +483,7 @@ def clip_gradients(
             )
         raise ValueError("You have to specify either `clip_val` or `max_norm` to do gradient clipping!")
 
-    def autocast(self) -> ContextManager:
+    def autocast(self) -> AbstractContextManager:
         """A context manager to automatically convert operations for the chosen precision.
 
         Use this only if the `forward` method of your model does not cover all operations you wish to run with the
@@ -634,7 +633,7 @@ def rank_zero_first(self, local: bool = False) -> Generator:
             if rank == 0:
                 barrier()
 
-    def no_backward_sync(self, module: _FabricModule, enabled: bool = True) -> ContextManager:
+    def no_backward_sync(self, module: _FabricModule, enabled: bool = True) -> AbstractContextManager:
         r"""Skip gradient synchronization during backward to avoid redundant communication overhead.
 
         Use this context manager when performing gradient accumulation to speed up training with multiple devices.
@@ -676,7 +675,7 @@ def no_backward_sync(self, module: _FabricModule, enabled: bool = True) -> Conte
         forward_module, _ = _unwrap_compiled(module._forward_module)
         return self._strategy._backward_sync_control.no_backward_sync(forward_module, enabled)
 
-    def sharded_model(self) -> ContextManager:
+    def sharded_model(self) -> AbstractContextManager:
         r"""Instantiate a model under this context manager to prepare it for model-parallel sharding.
 
         .. deprecated:: This context manager is deprecated in favor of :meth:`init_module`, use it instead.
@@ -688,12 +687,12 @@ def sharded_model(self) -> ContextManager:
             return self.strategy.module_sharded_context()
         return nullcontext()
 
-    def init_tensor(self) -> ContextManager:
+    def init_tensor(self) -> AbstractContextManager:
         """Tensors that you instantiate under this context manager will be created on the device right away and have
         the right data type depending on the precision setting in Fabric."""
         return self._strategy.tensor_init_context()
 
-    def init_module(self, empty_init: Optional[bool] = None) -> ContextManager:
+    def init_module(self, empty_init: Optional[bool] = None) -> AbstractContextManager:
         """Instantiate the model and its parameters under this context manager to reduce peak memory usage.
 
         The parameters get created on the device and with the right data type right away without wasting memory being

src/lightning/fabric/plugins/precision/amp.py

@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, ContextManager, Literal, Optional
+from contextlib import AbstractContextManager
+from typing import Any, Literal, Optional
 
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
@@ -59,7 +60,7 @@ def __init__(
         self._desired_input_dtype = torch.bfloat16 if self.precision == "bf16-mixed" else torch.float16
 
     @override
-    def forward_context(self) -> ContextManager:
+    def forward_context(self) -> AbstractContextManager:
         return torch.autocast(self.device, dtype=self._desired_input_dtype)
 
     @override

src/lightning/fabric/plugins/precision/bitsandbytes.py

@@ -17,10 +17,10 @@
 import os
 import warnings
 from collections import OrderedDict
-from contextlib import ExitStack
+from contextlib import AbstractContextManager, ExitStack
 from functools import partial
 from types import ModuleType
-from typing import Any, Callable, ContextManager, Literal, Optional, cast
+from typing import Any, Callable, Literal, Optional, cast
 
 import torch
 from lightning_utilities import apply_to_collection
@@ -123,11 +123,11 @@ def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
         return module
 
     @override
-    def tensor_init_context(self) -> ContextManager:
+    def tensor_init_context(self) -> AbstractContextManager:
         return _DtypeContextManager(self.dtype)
 
     @override
-    def module_init_context(self) -> ContextManager:
+    def module_init_context(self) -> AbstractContextManager:
         if self.ignore_modules:
             # cannot patch the Linear class if the user wants to skip some submodules
             raise RuntimeError(
@@ -145,7 +145,7 @@ def module_init_context(self) -> ContextManager:
         return stack
 
     @override
-    def forward_context(self) -> ContextManager:
+    def forward_context(self) -> AbstractContextManager:
         return _DtypeContextManager(self.dtype)
 
     @override

src/lightning/fabric/plugins/precision/deepspeed.py

@@ -11,8 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from contextlib import nullcontext
-from typing import TYPE_CHECKING, Any, ContextManager, Literal
+from contextlib import AbstractContextManager, nullcontext
+from typing import TYPE_CHECKING, Any, Literal
 
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
@@ -68,13 +68,13 @@ def convert_module(self, module: Module) -> Module:
         return module
 
     @override
-    def tensor_init_context(self) -> ContextManager:
+    def tensor_init_context(self) -> AbstractContextManager:
         if "true" not in self.precision:
             return nullcontext()
         return _DtypeContextManager(self._desired_dtype)
 
     @override
-    def module_init_context(self) -> ContextManager:
+    def module_init_context(self) -> AbstractContextManager:
         return self.tensor_init_context()
 
     @override

src/lightning/fabric/plugins/precision/double.py

@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, ContextManager, Literal
+from contextlib import AbstractContextManager
+from typing import Any, Literal
 
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
@@ -33,15 +34,15 @@ def convert_module(self, module: Module) -> Module:
         return module.double()
 
     @override
-    def tensor_init_context(self) -> ContextManager:
+    def tensor_init_context(self) -> AbstractContextManager:
         return _DtypeContextManager(torch.double)
 
     @override
-    def module_init_context(self) -> ContextManager:
+    def module_init_context(self) -> AbstractContextManager:
         return self.tensor_init_context()
 
     @override
-    def forward_context(self) -> ContextManager:
+    def forward_context(self) -> AbstractContextManager:
         return self.tensor_init_context()
 
     @override

src/lightning/fabric/plugins/precision/fsdp.py

@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import TYPE_CHECKING, Any, ContextManager, Literal, Optional
+from contextlib import AbstractContextManager
+from typing import TYPE_CHECKING, Any, Literal, Optional
 
 import torch
 from lightning_utilities import apply_to_collection
@@ -100,15 +101,15 @@ def mixed_precision_config(self) -> "TorchMixedPrecision":
         )
 
     @override
-    def tensor_init_context(self) -> ContextManager:
+    def tensor_init_context(self) -> AbstractContextManager:
         return _DtypeContextManager(self._desired_input_dtype)
 
     @override
-    def module_init_context(self) -> ContextManager:
+    def module_init_context(self) -> AbstractContextManager:
         return _DtypeContextManager(self.mixed_precision_config.param_dtype or torch.float32)
 
     @override
-    def forward_context(self) -> ContextManager:
+    def forward_context(self) -> AbstractContextManager:
         if "mixed" in self.precision:
             return torch.autocast("cuda", dtype=(torch.bfloat16 if self.precision == "bf16-mixed" else torch.float16))
         return self.tensor_init_context()

src/lightning/fabric/plugins/precision/half.py

@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, ContextManager, Literal
+from contextlib import AbstractContextManager
+from typing import Any, Literal
 
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
@@ -42,15 +43,15 @@ def convert_module(self, module: Module) -> Module:
         return module.to(dtype=self._desired_input_dtype)
 
     @override
-    def tensor_init_context(self) -> ContextManager:
+    def tensor_init_context(self) -> AbstractContextManager:
         return _DtypeContextManager(self._desired_input_dtype)
 
     @override
-    def module_init_context(self) -> ContextManager:
+    def module_init_context(self) -> AbstractContextManager:
         return self.tensor_init_context()
 
     @override
-    def forward_context(self) -> ContextManager:
+    def forward_context(self) -> AbstractContextManager:
         return self.tensor_init_context()
 
     @override

src/lightning/fabric/plugins/precision/precision.py

@@ -11,8 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from contextlib import nullcontext
-from typing import Any, ContextManager, Literal, Optional, Union
+from contextlib import AbstractContextManager, nullcontext
+from typing import Any, Literal, Optional, Union
 
 from torch import Tensor
 from torch.nn import Module
@@ -53,19 +53,19 @@ def convert_module(self, module: Module) -> Module:
         """
         return module
 
-    def tensor_init_context(self) -> ContextManager:
+    def tensor_init_context(self) -> AbstractContextManager:
         """Controls how tensors get created (device, dtype)."""
         return nullcontext()
 
-    def module_init_context(self) -> ContextManager:
+    def module_init_context(self) -> AbstractContextManager:
         """Instantiate module parameters or tensors in the precision type this plugin handles.
 
         This is optional and depends on the precision limitations during optimization.
 
         """
         return nullcontext()
 
-    def forward_context(self) -> ContextManager:
+    def forward_context(self) -> AbstractContextManager:
         """A contextmanager for managing model forward/training_step/evaluation_step/predict_step."""
         return nullcontext()
 

src/lightning/fabric/plugins/precision/transformer_engine.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 import logging
 from collections.abc import Mapping
-from contextlib import ExitStack
-from typing import TYPE_CHECKING, Any, ContextManager, Literal, Optional, Union
+from contextlib import AbstractContextManager, ExitStack
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 import torch
 from lightning_utilities import apply_to_collection
@@ -107,11 +107,11 @@ def convert_module(self, module: torch.nn.Module) -> torch.nn.Module:
         return module
 
     @override
-    def tensor_init_context(self) -> ContextManager:
+    def tensor_init_context(self) -> AbstractContextManager:
         return _DtypeContextManager(self.weights_dtype)
 
     @override
-    def module_init_context(self) -> ContextManager:
+    def module_init_context(self) -> AbstractContextManager:
         dtype_ctx = self.tensor_init_context()
         stack = ExitStack()
         if self.replace_layers:
@@ -126,7 +126,7 @@ def module_init_context(self) -> ContextManager:
         return stack
 
     @override
-    def forward_context(self) -> ContextManager:
+    def forward_context(self) -> AbstractContextManager:
         dtype_ctx = _DtypeContextManager(self.weights_dtype)
         fallback_autocast_ctx = torch.autocast(device_type="cuda", dtype=self.fallback_compute_dtype)
         import transformer_engine.pytorch as te

src/lightning/fabric/strategies/ddp.py

@@ -11,9 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from contextlib import nullcontext
+from contextlib import AbstractContextManager, nullcontext
 from datetime import timedelta
-from typing import Any, ContextManager, Literal, Optional, Union
+from typing import Any, Literal, Optional, Union
 
 import torch
 import torch.distributed
@@ -231,7 +231,7 @@ def _determine_ddp_device_ids(self) -> Optional[list[int]]:
 
 class _DDPBackwardSyncControl(_BackwardSyncControl):
     @override
-    def no_backward_sync(self, module: Module, enabled: bool) -> ContextManager:
+    def no_backward_sync(self, module: Module, enabled: bool) -> AbstractContextManager:
         """Blocks gradient synchronization inside the :class:`~torch.nn.parallel.distributed.DistributedDataParallel`
         wrapper."""
         if not enabled: