Fix: Revert lightning_lite.utilities.rank_zero_only to preserve backward compatibility (#15536)

* Fix: Revert  to preserve backward compatibility

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/pytorch_lightning/CHANGELOG.md

@@ -12,6 +12,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added utilities to migrate checkpoints from one Lightning version to another ([#15237](https://github.com/Lightning-AI/lightning/pull/15237))
 
+- Added back the accidentally removed `pytorch_lightning.utilities.distributed.rank_zero_only` function ([#15536](https://github.com/Lightning-AI/lightning/pull/15536))
+
 - Added support to upgrade all checkpoints in a folder using the `pl.utilities.upgrade_checkpoint` script ([#15333](https://github.com/Lightning-AI/lightning/pull/15333))
 
 -
@@ -28,7 +30,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Deprecated
 
--
+- Deprecated `pytorch_lightning.utilities.distributed.rank_zero_only` in favor of `pytorch_lightning.utilities.rank_zero_only` ([#15536](https://github.com/Lightning-AI/lightning/pull/15536))
 
 -
 

src/pytorch_lightning/utilities/distributed.py

@@ -252,3 +252,13 @@ def tpu_distributed() -> bool:
     from lightning_lite.accelerators.tpu import _tpu_distributed
 
     return _tpu_distributed()
+
+
+def rank_zero_only(*args: Any, **kwargs: Any) -> Any:
+    rank_zero_deprecation(
+        "`pytorch_lightning.utilities.distributed.rank_zero_only` has been deprecated in v1.8.1 and will"
+        " be removed in v1.10.0. You can import it from `pytorch_lightning.utilities` instead."
+    )
+    from pytorch_lightning.utilities.rank_zero import rank_zero_only as new_rank_zero_only
+
+    return new_rank_zero_only(*args, **kwargs)

src/pytorch_lightning/utilities/rank_zero.py

@@ -15,7 +15,6 @@
 import logging
 
 # note: we want to keep these indirections so the `rank_zero_module.log` is set (on import) for PL users
-# backwards compatibility
 from lightning_lite.utilities.rank_zero import LightningDeprecationWarning  # noqa: F401
 from lightning_lite.utilities.rank_zero import (  # noqa: F401
     rank_zero_debug,

tests/tests_pytorch/deprecated_api/test_remove_1-10.py

@@ -332,3 +332,10 @@ def test_v1_8_0_deprecated_all_gather_grad():
 
         with pytest.deprecated_call(match="`AllGatherGrad` has been deprecated in v1.8"):
             AllGatherGrad.apply(tensor1)
+
+
+def test_v1_8_1_deprecated_rank_zero_only():
+    from pytorch_lightning.utilities.distributed import rank_zero_only
+
+    with pytest.deprecated_call(match="rank_zero_only` has been deprecated in v1.8.1"):
+        rank_zero_only(lambda: None)