File size: 2,270 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
# Copyright The Lightning AI team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
"""Warning-related utilities."""
import warnings
from pathlib import Path
from typing import Optional, Union
from lightning_fabric.utilities.rank_zero import LightningDeprecationWarning
# enable our warnings
warnings.simplefilter("default", category=LightningDeprecationWarning)
_default_format_warning = warnings.formatwarning
def disable_possible_user_warnings(module: str = "") -> None:
"""Ignore warnings of the category ``PossibleUserWarning`` from Lightning.
For more granular control over which warnings to ignore, use :func:`warnings.filterwarnings` directly.
Args:
module: Name of the module for which the warnings should be ignored (e.g., ``'pytorch_lightning.strategies'``).
Default: Disables warnings from all modules.
"""
warnings.filterwarnings("ignore", module=module, category=PossibleUserWarning)
def _custom_format_warning(
message: Union[Warning, str], category: type[Warning], filename: str, lineno: int, line: Optional[str] = None
) -> str:
"""Custom formatting that avoids an extra line in case warnings are emitted from the `rank_zero`-functions."""
if _is_path_in_lightning(Path(filename)):
# The warning originates from the Lightning package
return f"{filename}:{lineno}: {message}\n"
return _default_format_warning(message, category, filename, lineno, line)
warnings.formatwarning = _custom_format_warning
class PossibleUserWarning(UserWarning):
"""Warnings that could be false positives."""
def _is_path_in_lightning(path: Path) -> bool:
"""Naive check whether the path looks like a path from the lightning package."""
return "lightning" in str(path.absolute())
|